One problem with JAX is that it is
JAX has checkify, which enables error checking. The idea is that the compiled program returns a status code that can be checked afterwards.
The user interface is that a checkified function returns an error code and the actual return value and you can check it along the lines of:
f <- checkify(jit(f))
out <- f(...)
if (is_error(out[[1]])) {
stop("something went wrong")
} else {
print("all good, keep running")
}This is similar to how gomlx handles printing, i.e. every printed tensor is simply appended as an output to the program.
It's also reminiscent of LAPACK.
Via these status codes, things like bounds checks could be implemented.
I.e., within the kernel we could compute all((indices <= end) & (indices >= start)) and return this as an additional argument.
The trickiest thing to handle are higher order primitives and composing with other transformations.
It would be great if we could run such checks.
E.g. currently, probs vector in quantile is static because we otherwise can't check that they are within [0, 1].
The simplest solution would just be to await the return after calling a kernel, but this is obviously horrible for performance reasons.
The direction that seems the most promising is that every buffer keeps a chain of events that preceded it.
I think this chain must either be limited in length (and once the length is exceeded we insert a host sync) or we implement it via some linked list (if we use a dynamically sized vector and have very long kernel queues, this becomes a performance issue.
Having this would certainly be some work, but I think it would be much nicer for end users. I really dislike some of XLA's semantics that just clamp indices etc.
One problem with JAX is that it is
JAX has
checkify, which enables error checking. The idea is that the compiled program returns a status code that can be checked afterwards.The user interface is that a checkified function returns an error code and the actual return value and you can check it along the lines of:
This is similar to how gomlx handles printing, i.e. every printed tensor is simply appended as an output to the program.
It's also reminiscent of LAPACK.
Via these status codes, things like bounds checks could be implemented.
I.e., within the kernel we could compute
all((indices <= end) & (indices >= start))and return this as an additional argument.The trickiest thing to handle are higher order primitives and composing with other transformations.
It would be great if we could run such checks.
E.g. currently,
probsvector in quantile is static because we otherwise can't check that they are within[0, 1].The simplest solution would just be to await the return after calling a kernel, but this is obviously horrible for performance reasons.
The direction that seems the most promising is that every buffer keeps a chain of events that preceded it.
I think this chain must either be limited in length (and once the length is exceeded we insert a host sync) or we implement it via some linked list (if we use a dynamically sized vector and have very long kernel queues, this becomes a performance issue.
Having this would certainly be some work, but I think it would be much nicer for end users. I really dislike some of XLA's semantics that just clamp indices etc.