Add runtime_raise inside defn by Chapaman · Pull Request #1822 · elixir-nx/nx

Chapaman · 2026-08-19T23:58:43Z

Nx.Defn.Kernel.runtime_raise inside if/cond/while runs at runtime when that branch is taken. raise stays a trace-time check for shapes, ranks, and case.

Failed EXLA callbacks used to crash the Runner GenServer. They now exit the Outfeed with CallbackError and raise in the original caller, so they can be rescued or tested with assert_raise.

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.

polvalente · 2026-08-20T13:06:09Z

+  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


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.

Went with throw-during-trace, catch in cond/while, then the existing io_call raise shape. No new Expr op: those are tensors and raise has no type/shape. case stays trace-time (it's Elixir case). raise_if is now if + raise.

polvalente · 2026-08-20T13:08:17Z

      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)


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

run no longer unwrap-throws. Failed status comes back as {:error, message} from the NIF. Executable.run still raises that so the public contract is unchanged. Callback exceptions themselves come from the Outfeed DOWN.

josevalim · 2026-08-21T09:45:01Z

+
+        callback_error =
+          receive do
+            {:exla_callback_error, ^error_ref, kind, reason, stacktrace} ->


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.

Done. Outfeed exits with {:shutdown, %EXLA.Defn.CallbackError{...}} and maybe_outfeed unpacks the DOWN reason. {:shutdown, _} so Task.Supervisor doesn't log it. No error ref, extra messages, or passing self.

Outfeed exits with CallbackError instead of sending a side message, and the NIF returns {:error, message} instead of throwing on a failed Run.

Trace-time raise still applies at the top level and in case. raise_if is now sugar over if plus raise.

polvalente · 2026-08-24T22:46:20Z

  defmacro raise(exception, arguments) do
    quote do
-      Elixir.Kernel.raise(unquote(exception), unquote(arguments))
+      Nx.Defn.Expr.trace_raise(unquote(exception), unquote(arguments))
    end


I think if we denormalize this node to regular Elixir.Kernel.raise when inside :case, we can implement Nx.Defn.Expr.raise such that it is always runtime only.

https://github.com/Chapaman/nx/blob/0fe88be090643087b1ce03a836b8289bc4900ba7/nx/lib/nx/defn/compiler.ex#L485

If this works, we can then remove raise-if and the throw-catch trick

Case already ends up as Kernel.raise: only the matching clause runs, the throw escapes, and the compiler converts it. Rewriting in normalize({:case}) is the same outcome.

Always-runtime Expr.raise doesn't cover if Nx.rank(t) != 2, do: raise then {r, c} = shape. Rank is only known while tracing, and a discarded runtime node lets tracing continue into the match (that's the LinAlg CI failure). 4392b55 Kernel.raises when the remaining if/cond clause is already known true.

I couldn't drop the throw-catch: if raise only returns a node, raise; value falls through. Same question for raise_if as the pipeable wrapper. Do you have something in mind for those?

@polvalente ah, the issue is that we already allow raise while building the graph, and that actually raises, so we cannot convert the same raise to raise at runtime. We will likely need runtime_raise or io_raise or something instead. :( Thoughts?

I like runtime_raise!

Went with runtime_raise. raise is back to Kernel.raise while building the graph. raise_if is if plus runtime_raise.

Let's remove raise_if. The original reason for adding it was for disambiguation of plain raise

A discarded compile-time raise was turning into an io_call, so tracing continued into later pattern matches and LinAlg doctests got MatchError.

raise already meant "while building the graph", so it cannot also be a runtime host callback. runtime_raise is that callback; raise_if is if plus runtime_raise.

polvalente · 2026-08-26T17:09:10Z

+    defn custom_runtime_raise(value, predicate) do
+      if predicate do
+        runtime_raise RuntimeRaiseError,
+          message: "custom runtime check failed",
+          value: :preserved
+      else
+        value
+      end
+    end
+


we might want a test where we have a custom call raise and catch inside itself to ensure that we only bubble up uncaught raises

polvalente · 2026-08-26T17:10:04Z

+        while {x, i = 0, n}, i < 10 do
+          i =
+            if i == n do
+              runtime_raise "Halting on selected iteration"


Maybe have one of these include the iteration as an argument? Not sure if the tensor would survive allocation, though. Seems like an interesting experiment.

the tensor does survive when passing through io_call

polvalente · 2026-08-26T17:11:44Z

+      case Nx.Defn.Expr.catch_runtime_raise(fn ->
+             fun
+             |> apply(args)
+             |> Nx.Defn.Composite.traverse(&Nx.Defn.Expr.tensor/1)
+           end) do
+        {:value, result} ->
+          result
+
+        {:raise, spec} ->
+          Nx.Defn.Expr.apply_runtime_raise(spec)
+      end
    after
      if previous_backend do
        Process.put(Nx.Shared.backend_pdict_key(), previous_backend)


Do we still need this throw-catch mechanism given that runtime_raise is a new Expr node?
I believe we can have the compiler introduce the raise as a runtime thing when it is reached during execution.

after our talk, i'll have a ghost tensor to flag so that it is included on the result but removed before the final return - inspired by what is done on callback pid on exla

I think this is due to how tracing is implemented. runtime_raise should be functionally equivalent to an io_call that raises, but it has a special node to not require the reassignment and passing a tensor as argument

It was only there to disambiguate from raise. Also cover a rescued io_call callback and raising with the runtime loop index.

The branch result is passed through the callback, like other io_call values, so the raise stays in the graph without a tracing throw.

josevalim · 2026-08-27T07:00:42Z

+      Outfeed.start_child(
+        executable,
+        outfeed,
+        Process.group_leader(),
+        infeeds
+      )


We can revert this change!

josevalim · 2026-08-27T07:01:44Z

+        try do
+          EXLA.Executable.run(executable, [Enum.reverse(buffers)], run_options)
+        after
+          send(outfeed_pid, :stop)


Why is this necessary? Shouldn't we keep it in the outfeed as before?

josevalim · 2026-08-27T07:02:34Z

+          {{:shutdown, %EXLA.Defn.CallbackError{} = error}, _runner_result} ->
+            :erlang.raise(error.kind, error.reason, error.stacktrace)
+
+          {_down_reason, {:error, kind, reason, stacktrace}} ->
+            :erlang.raise(kind, reason, stacktrace)


I don't understand why we have both the CallbackError and the {:error, kind, reason, stacktrace}. Could they be unified into one?

josevalim · 2026-08-27T07:03:21Z

-      decompose_output(data_and_device_id, output_typespecs, client, mesh)
+    case run(client, ref, device_id, inputs, options) do
+      {:error, message} ->
+        raise RuntimeError, message: IO.iodata_to_binary(message)


Let's return {:error, message} to the caller and let the caller choose to raise.

josevalim · 2026-08-27T07:03:41Z

+      results ->
+        for data_and_device_id <- results do
+          decompose_output(data_and_device_id, output_typespecs, client, mesh)
+        end


Let's make this proper tagged tuples:

Suggested change

results ->

for data_and_device_id <- results do

decompose_output(data_and_device_id, output_typespecs, client, mesh)

end

{:ok, results} ->

{:ok, for data_and_device_id <- results do

decompose_output(data_and_device_id, output_typespecs, client, mesh)

end}

josevalim · 2026-08-27T07:07:13Z

I have dropped some comments but it seems the implementation of Nx.Defn.Expr is still a bit too complex. Could runtime_raise be equivalent to:

defmacro runtime_raise(arg) do
  quote do
    io_call(fn -> Kernel.raise(arg) end)
  end
end

or similar?

On the other hand, we may have issues with tensors and inspect and what not, so an alternative is:

defmacro runtime_raise(arg) when is_binary(arg) do
  quote do
    io_call(fn -> Kernel.raise(arg) end)
  end
end

where we require raise to be a compile-time string.

Chapaman · 2026-08-28T23:53:53Z

I have dropped some comments but it seems the implementation of Nx.Defn.Expr is still a bit too complex. Could runtime_raise be equivalent to:
defmacro runtime_raise(arg) do
  quote do
    io_call(fn -> Kernel.raise(arg) end)
  end
end
or similar?

On the other hand, we may have issues with tensors and inspect and what not, so an alternative is:
defmacro runtime_raise(arg) when is_binary(arg) do
  quote do
    io_call(fn -> Kernel.raise(arg) end)
  end
end
where we require raise to be a compile-time string.

Dropped the Expr tracing stuff. runtime_raise is now io_call(Nx.tensor(0), fn _ -> Kernel.raise(msg) end). The message is a compile-time string to avoid dealing with tensors/inspect, is this ok?

Chapaman added 2 commits August 19, 2026 20:58

Raise EXLA callback errors in the caller

770f107

Keep callback failures structured on the BEAM so they halt execution without crashing the linked runner process.

Add runtime raise_if to Nx.Defn.Kernel

ed26a11

Allow defn computations to halt from tensor predicates while preserving the checked value on the non-raising path.

Chapaman changed the title ~~Raise EXLA callback errors in the caller~~ Add runtime raise_if to Nx.Defn.Kernel Aug 20, 2026

Chapaman marked this pull request as ready for review August 20, 2026 01:18

polvalente reviewed Aug 20, 2026

View reviewed changes

josevalim reviewed Aug 21, 2026

View reviewed changes

Chapaman added 2 commits August 24, 2026 19:32

Surface EXLA callback errors via the Outfeed DOWN reason

c24ea2a

Outfeed exits with CallbackError instead of sending a side message, and the NIF returns {:error, message} instead of throwing on a failed Run.

Make raise work at runtime inside if, cond, and while

0fe88be

Trace-time raise still applies at the top level and in case. raise_if is now sugar over if plus raise.

polvalente reviewed Aug 24, 2026

View reviewed changes

Raise at trace time when an if/cond predicate is already known

4392b55

A discarded compile-time raise was turning into an io_call, so tracing continued into later pattern matches and LinAlg doctests got MatchError.

Chapaman changed the title ~~Add runtime raise_if to Nx.Defn.Kernel~~ Make raise work at runtime inside defn Aug 25, 2026

Keep raise at trace time and add runtime_raise

239a11c

raise already meant "while building the graph", so it cannot also be a runtime host callback. runtime_raise is that callback; raise_if is if plus runtime_raise.

Chapaman changed the title ~~Make raise work at runtime inside defn~~ Add runtime_raise inside defn Aug 26, 2026

polvalente reviewed Aug 26, 2026

View reviewed changes

Chapaman added 3 commits August 26, 2026 17:37

Drop raise_if now that runtime_raise exists

397f2bf

It was only there to disambiguate from raise. Also cover a rescued io_call callback and raising with the runtime loop index.

Thread runtime_raise through io_call instead of throwing

2a36fc9

The branch result is passed through the callback, like other io_call values, so the raise stays in the graph without a tracing throw.

Add parens to runtime_raise in tests for Elixir 1.20 formatting.

8476745

josevalim reviewed Aug 27, 2026

View reviewed changes

Chapaman added 2 commits August 28, 2026 20:42

Make runtime_raise a thin io_call instead of tracing machinery.

02f2062

Return tagged tuples from Executable.run and drop the runner catch.

c29efc1

Uh oh!

Conversation

Chapaman commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Chapaman Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Chapaman Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

josevalim commented Aug 27, 2026

Uh oh!

Chapaman commented Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Chapaman commented Aug 19, 2026 •

edited

Loading

Chapaman Aug 24, 2026 •

edited

Loading

Chapaman Aug 26, 2026 •

edited

Loading