From 74ebcab8e4fd2775751c15c8ae7b2a199d3921fb Mon Sep 17 00:00:00 2001 From: Bradley Lewis Fargo Date: Tue, 20 Oct 2026 10:25:31 -0500 Subject: [PATCH 1/2] Give block bodies their own expression context A block body is compiled as a function of the block's arguments, so a tensor it reads from the enclosing scope cannot be resolved there. The Evaluator returned an unevaluated expression, EXLA failed to find the parameter, and the gradient of such a tensor came out as zeros. The block parameters shared the caller's context, so nothing caught this. Trace the body under a context of its own, as while does, so an operation that mixes a parameter with an outer tensor raises when the block is built. A body that returns an outer tensor unchanged is checked separately. Co-Authored-By: Claude Fable 5.1 --- nx/CHANGELOG.md | 1 + nx/lib/nx/defn/expr.ex | 30 +++++++++++++++++++++++++----- nx/test/nx/defn_test.exs | 28 ++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 5 deletions(-) diff --git a/nx/CHANGELOG.md b/nx/CHANGELOG.md index 3a592a0600..6ab3cc92e1 100644 --- a/nx/CHANGELOG.md +++ b/nx/CHANGELOG.md @@ -8,6 +8,7 @@ ### Bug fixes + * [Nx.Defn] Raise when the function given to a block reads a tensor from the enclosing scope instead of returning wrong results * [Nx.Defn] Apply gradients for operations that create constants inside blocks * [Nx.Defn] Preserve tensor dimensions that match vectorized axis sizes in conditional gradients * [Nx.Defn] Avoid counting gradient contributions twice when differentiating conditionals diff --git a/nx/lib/nx/defn/expr.ex b/nx/lib/nx/defn/expr.ex index fd44b273fd..def97138e0 100644 --- a/nx/lib/nx/defn/expr.ex +++ b/nx/lib/nx/defn/expr.ex @@ -836,8 +836,8 @@ defmodule Nx.Defn.Expr do @impl true def block(struct, _output \\ nil, in_args, fun) do {args, opts} = Enum.split_while(in_args, &(not is_list(&1))) - {args, context} = to_exprs(args) - context = context || :root + {args, parent_context} = to_exprs(args) + context = parent_context || :root # Nodes store devectorized tensors, so the body is traced on the # vectorized view and the result is vectorized again on the way out. @@ -845,22 +845,42 @@ defmodule Nx.Defn.Expr do devectorized_args = Enum.map(args, &Nx.devectorize/1) in_args = devectorized_args ++ opts + # The body gets its own context, so an operation that mixes a parameter + # with a tensor from the enclosing scope fails when the block is built. + # Compilers run the body as a function of its arguments alone and cannot + # reach such a tensor. + body_context = new_context(:block) + params = - Enum.with_index(devectorized_args, fn arg, pos -> parameter(arg, context, pos) end) + Enum.with_index(devectorized_args, fn arg, pos -> parameter(arg, body_context, pos) end) vectorized_params = Enum.zip_with(params, vectorized_axes, &vectorize_axes/2) traced = apply(fun, [struct | vectorized_params ++ opts]) + + # A tensor from the enclosing scope returned as is would bypass the check + # above. Outside a trace the arguments have no context to compare with. + if parent_context != nil do + Composite.reduce(traced, nil, fn %{data: %{context: traced_context}}, acc -> + if traced_context == parent_context do + raise ArgumentError, + "the function given to Nx.block returned a tensor from the enclosing scope. " <> + "Every tensor the function uses must be given as one of the block arguments" + end + + acc + end) + end + out_vectorized_axes = [traced] |> Composite.flatten_list() |> Enum.map(& &1.vectorized_axes) fun = vectorized_block_fun(fun, vectorized_axes, length(opts)) case Composite.traverse(traced, &Nx.devectorize/1) do - %{data: %{context: context}} = res -> + %{} = res -> res |> expr(context, :block, [struct, in_args, res, fun]) |> revectorize_block_result(out_vectorized_axes) t when is_tuple(t) -> - context = elem(t, 0).data.context out = tuple_out(tuple_size(t)) node = expr(out, context, :block, [struct, in_args, t, fun]) revectorize_block_result(tuple(node, Tuple.to_list(t)), out_vectorized_axes) diff --git a/nx/test/nx/defn_test.exs b/nx/test/nx/defn_test.exs index 4fe9e6f4ad..99a954199d 100644 --- a/nx/test/nx/defn_test.exs +++ b/nx/test/nx/defn_test.exs @@ -937,6 +937,34 @@ defmodule Nx.DefnTest do defstruct [] end + test "raises when the body mixes an argument with a tensor from the enclosing scope" do + assert_raise RuntimeError, ~r/expressions come from different contexts/, fn -> + Nx.Defn.jit(fn x, weight -> + Nx.block(%VectorizedNonlinear{}, [x], nil, fn _, x -> Nx.multiply(x, weight) end) + end).(Nx.tensor([1.0, 2.0, 3.0]), Nx.tensor([2.0, 3.0, 4.0])) + end + end + + test "raises when the body returns a tensor from the enclosing scope" do + assert_raise ArgumentError, ~r/returned a tensor from the enclosing scope/, fn -> + Nx.Defn.jit(fn x, y -> + Nx.block(%VectorizedNonlinear{}, [x], nil, fn _, _x -> y end) + end).(Nx.tensor([1.0, 2.0]), Nx.tensor([3.0, 4.0])) + end + end + + test "raises when the body reads a tensor from the enclosing scope inside grad" do + assert_raise RuntimeError, ~r/expressions come from different contexts/, fn -> + Nx.Defn.jit(fn x, w -> + grad(w, fn w -> + Nx.sum( + Nx.block(%VectorizedNonlinear{}, [x], nil, fn _, x -> Nx.sin(Nx.dot(x, w)) end) + ) + end) + end).(Nx.tensor([[0.5, 1.0], [1.5, 2.0]]), Nx.tensor([[0.5, -0.3], [0.2, 0.8]])) + end + end + @tag compiler: Evaluator test "applies gradients for a block whose output is a constant-creating comparison" do input = Nx.tensor([-0.7, 0.3, 0.7]) From da449777196097946e1d58ffd14d498c264847d2 Mon Sep 17 00:00:00 2001 From: Bradley Lewis Fargo Date: Tue, 20 Oct 2026 11:18:31 -0500 Subject: [PATCH 2/2] Rerun CI Co-Authored-By: Claude Fable 5.1