Add runtime raise_if to Nx.Defn.Kernel - #1822
Conversation
Keep callback failures structured on the BEAM so they halt execution without crashing the linked runner process.
Allow defn computations to halt from tensor predicates while preserving the checked value on the non-raising path.
Chapaman
changed the title
| defmacro raise_if(value, predicate, exception_or_message) do | ||
| quote do | ||
| value = unquote(value) | ||
|
|
||
| Nx.Defn.Kernel.if(unquote(predicate), | ||
| do: | ||
| Nx.Defn.Kernel.io_call(value, fn _ -> | ||
| Elixir.Kernel.raise(unquote(exception_or_message)) | ||
| end), | ||
| else: value | ||
| ) | ||
| end | ||
| end |
There was a problem hiding this comment.
I think we should try to make raise itself become a special node that works in runtime when called inside defn. I'm just not sure how to make raise work inside case. raise_if is a nice out if we can't make it happen, though.
| EXLA.Defn.Runner.start_link(lock, fn -> | ||
| EXLA.Executable.run(executable, [Enum.reverse(buffers)], run_options) | ||
| try do | ||
| EXLA.Executable.run(executable, [Enum.reverse(buffers)], run_options) |
There was a problem hiding this comment.
I think we're missing a change in the c++ code to not raise when unwrapping the result. This way if the callback returns {:error, ...} we can just surface that tuple without raising
|
|
||
| callback_error = | ||
| receive do | ||
| {:exla_callback_error, ^error_ref, kind, reason, stacktrace} -> |
There was a problem hiding this comment.
It would be nice to not have to deal with additional messages. One option is to wrap the raised exception in EXLA.Defn.CallbackError{kind: kind, reason: reason, stacktrace: stacktrace} and then in here, you can check the DOWN reason (the fifth element of the tuple) and unpack it accordingly. Then we don't need the error ref, we don't need to pass self, etc.
3 participants
Allow
Nx.Defn.Kernel.raise_if/3andraise_if/4to halt a computation from a runtime tensor predicate.Previously, a failed EXLA callback crashed the Runner GenServer and surfaced as a process exit. Callback failures now stop the computation but raise normally in the original caller, where they can be rescued or tested with
assert_raise.raise_ifreturns the checked value unchanged when its predicate is false and supports both messages and custom exceptions.