fix: Nx.block with EXLA JIT (#1750)

Author: polvalente

exla/test/exla/defn/api_test.exs

@@ -320,6 +320,42 @@ defmodule EXLA.Defn.APITest do
     end
   end
 
+  describe "block ops via Evaluator with EXLA.Backend" do
+    # Regression test: before the fix, eval_apply(:block, ...) wrapped block_apply_default
+    # in a Nx.Defn.Compiler.fun and passed it to backend.block. EXLA.Backend.block JIT-traces
+    # that fun through Nx.Defn.Expr.block, calling it with Expr params. block_apply_default
+    # then hit composite_eval which raised on Expr tensors at the :parameter guard.
+    #
+    # The fix: pass the original callback directly so JIT-tracing works as intended.
+    setup do
+      Nx.default_backend({EXLA.Backend, client: :host})
+      # Override the global EXLA compiler set in test_helper so Evaluator handles defn dispatch
+      Nx.Defn.default_options(compiler: Nx.Defn.Evaluator)
+      :ok
+    end
+
+    defn block_top_k(t), do: Nx.top_k(t, k: 2)
+    defn block_cumulative_sum(t), do: Nx.cumulative_sum(t)
+    defn block_logical_not(t), do: Nx.logical_not(t)
+
+    test "top_k" do
+      t = Nx.tensor([3.0, 1.0, 4.0, 1.0, 5.0], backend: {EXLA.Backend, client: :host})
+      {values, indices} = block_top_k(t)
+      assert_equal(values, Nx.tensor([5.0, 4.0]))
+      assert_equal(indices, Nx.tensor([4, 2]))
+    end
+
+    test "cumulative_sum" do
+      t = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: {EXLA.Backend, client: :host})
+      assert_equal(block_cumulative_sum(t), Nx.tensor([1.0, 3.0, 6.0, 10.0]))
+    end
+
+    test "logical_not" do
+      t = Nx.tensor([1, 0, 1, 0], backend: {EXLA.Backend, client: :host})
+      assert_equal(block_logical_not(t), Nx.tensor([0, 1, 0, 1], type: :u8))
+    end
+  end
+
   describe "telemetry" do
     defn telemetry_add_two(a, b), do: a + b
 

nx/lib/nx/defn/evaluator.ex

@@ -134,25 +134,13 @@ defmodule Nx.Defn.Evaluator do
   end
 
   defp compute_cache(:block, %{data: %Expr{args: args}}, state, cache) do
-    [struct, in_args, expr, _callback] = args
-    %module{} = struct
+    [struct, in_args, expr, callback] = args
 
     {call_prefix, call_suffix} = Enum.split_while(in_args, &(not is_list(&1)))
     {call_prefix, cache} = Enum.map_reduce(call_prefix, cache, &compute_cache(&1, state, &2))
     in_args = call_prefix ++ call_suffix
-    key = computation_key(module, call_prefix)
 
-    {{expr, expr_cache}, cache} =
-      case cache do
-        %{^key => optional_expr_cache} ->
-          {optional_expr_cache, cache}
-
-        %{} ->
-          optional_expr_cache = init_compute_cache(expr, state)
-          {optional_expr_cache, Map.put(cache, key, optional_expr_cache)}
-      end
-
-    {[struct, in_args, expr, expr_cache], cache}
+    {[struct, in_args, expr, callback], cache}
   end
 
   defp compute_cache(:cond, %{data: %Expr{args: [clauses, last]}}, state, cache) do
@@ -229,16 +217,6 @@ defmodule Nx.Defn.Evaluator do
     Tree.apply_args(tensor, cache, &compute_cache(&1, state, &2))
   end
 
-  defp computation_key(op, args) do
-    keys =
-      Enum.map(args, fn
-        %Nx.Tensor{shape: shape, names: names, type: type} -> {type, shape, names}
-        opts -> opts
-      end)
-
-    {op, keys}
-  end
-
   ## Evaluation
 
   defp eval(%Nx.Tensor{data: %Expr{op: :tensor, args: [t]}}, _state, caches) do
@@ -365,7 +343,7 @@ defmodule Nx.Defn.Evaluator do
     {{}, caches}
   end
 
-  defp eval_apply(:block, [struct, in_args, expr, expr_cache], ans, state, caches) do
+  defp eval_apply(:block, [struct, in_args, expr, callback], ans, state, caches) do
     {in_args, caches} = Enum.map_reduce(in_args, caches, &eval(&1, state, &2))
     {param_prefix, _} = Enum.split_while(in_args, &(not is_list(&1)))
     backend = Nx.Shared.list_impl!(param_prefix)
@@ -376,16 +354,7 @@ defmodule Nx.Defn.Evaluator do
         _ -> ans
       end
 
-    fun =
-      Nx.Defn.Compiler.fun(
-        length(in_args) + 1,
-        fn args ->
-          [struct | tensors] = args
-          block_apply_default(expr, state, expr_cache, struct, tensors)
-        end
-      )
-
-    {backend.block(struct, out, in_args, fun), caches}
+    {backend.block(struct, out, in_args, callback), caches}
   end
 
   defp eval_apply(:runtime_call, [expr, fun, out_template, opts], _ans, state, caches) do
@@ -429,11 +398,6 @@ defmodule Nx.Defn.Evaluator do
     {apply(mod, op, args), caches}
   end
 
-  defp block_apply_default(expr, state, expr_cache, _struct, args) when is_list(args) do
-    params = Enum.map(args, &fn -> &1 end)
-    elem(composite_eval(expr, %{state | params: params}, [expr_cache]), 0)
-  end
-
   ## Control flow helpers
 
   defp while(acc, condition, block, state, caches) do