Tiny fixes on block

Author: josevalim

exla/lib/exla/backend.ex

@@ -330,7 +330,7 @@ defmodule EXLA.Backend do
     {tensors, rest} = Enum.split_while(args, &is_struct(&1, Nx.Tensor))
 
     wrapper_fun = fn tensors ->
-      Nx.Defn.Expr.block(struct, nil, Tuple.to_list(tensors) ++ rest, fun)
+      Nx.Defn.Expr.block(struct, Tuple.to_list(tensors) ++ rest, fun)
     end
 
     jit([], wrapper_fun, tensors, [List.to_tuple(tensors)])

nx/lib/nx.ex

@@ -8935,7 +8935,7 @@ defmodule Nx do
     else
       out = %{a | names: [], shape: {}, type: {:u, 8}}
 
-      block(struct(Nx.Block.AllClose, opts), [a, b], out, fn %Nx.Block.AllClose{} = o, a, b ->
+      block(struct!(Nx.Block.AllClose, opts), [a, b], out, fn %Nx.Block.AllClose{} = o, a, b ->
         vectorized_all_close(a, b,
           equal_nan: o.equal_nan,
           rtol: o.rtol,
@@ -14327,19 +14327,18 @@ defmodule Nx do
       indices = devectorize(indices, keep_names: false)
       out = %{tensor | shape: inner_shape, names: inner_names}
 
-      block(struct(Nx.Block.Take, axis: axis), [tensor, indices], out, fn %Nx.Block.Take{},
-                                                                          tensor,
-                                                                          indices ->
-        gather_indices = new_axis(indices, rank(indices))
-        {indices_axes, tensor_axes} = Enum.split(axes(inner_shape), rank(indices))
-        {leading, trailing} = Enum.split(tensor_axes, axis)
+      block(struct!(Nx.Block.Take, axis: axis), [tensor, indices], out, fn
+        %Nx.Block.Take{}, tensor, indices ->
+          gather_indices = new_axis(indices, rank(indices))
+          {indices_axes, tensor_axes} = Enum.split(axes(inner_shape), rank(indices))
+          {leading, trailing} = Enum.split(tensor_axes, axis)
 
-        transpose_axes = leading ++ indices_axes ++ trailing
+          transpose_axes = leading ++ indices_axes ++ trailing
 
-        tensor
-        |> gather(gather_indices, axes: [axis])
-        |> transpose(axes: transpose_axes)
-        |> rename(inner_names)
+          tensor
+          |> gather(gather_indices, axes: [axis])
+          |> transpose(axes: transpose_axes)
+          |> rename(inner_names)
       end)
     end
   end
@@ -14509,7 +14508,7 @@ defmodule Nx do
 
     result =
       block(
-        struct(Nx.Block.TakeAlongAxis, axis: axis),
+        struct!(Nx.Block.TakeAlongAxis, axis: axis),
         [tensor, indices],
         out,
         fn %Nx.Block.TakeAlongAxis{}, tensor, indices ->
@@ -15327,7 +15326,7 @@ defmodule Nx do
       out_indices = %{tensor | shape: output_shape, names: output_names, type: {:s, 32}}
 
       block(
-        struct(Nx.Block.TopK, k: opts[:k]),
+        struct!(Nx.Block.TopK, k: opts[:k]),
         [tensor],
         {out_values, out_indices},
         fn %Nx.Block.TopK{} = top_k, tensor ->
@@ -16774,9 +16773,9 @@ defmodule Nx do
 
       block_struct =
         if kind == :fft2 do
-          struct(Nx.Block.FFT2, eps: opts[:eps], lengths: [l1, l2], axes: [ax1, ax2])
+          struct!(Nx.Block.FFT2, eps: opts[:eps], lengths: [l1, l2], axes: [ax1, ax2])
         else
-          struct(Nx.Block.IFFT2, eps: opts[:eps], lengths: [l1, l2], axes: [ax1, ax2])
+          struct!(Nx.Block.IFFT2, eps: opts[:eps], lengths: [l1, l2], axes: [ax1, ax2])
         end
 
       block(block_struct, [tensor], out, fn s, tensor ->

nx/lib/nx/defn/evaluator.ex

@@ -366,7 +366,7 @@ defmodule Nx.Defn.Evaluator do
   end
 
   defp eval_apply(:block, [struct, in_args, expr, expr_cache], ans, state, caches) do
-    {in_args, caches} = Tree.map_block_args(in_args, caches, &eval(&1, state, &2))
+    {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)
 

nx/lib/nx/defn/expr.ex

@@ -808,7 +808,7 @@ defmodule Nx.Defn.Expr do
   @behaviour Nx.Backend
 
   @impl true
-  def block(struct, _output, in_args, fun) do
+  def block(struct, _output \\ nil, in_args, fun) do
     expr_block(struct, in_args, fun)
   end
 

nx/lib/nx/defn/tree.ex

@@ -181,7 +181,7 @@ defmodule Nx.Defn.Tree do
 
   def apply_args(%T{data: %Expr{op: :block, args: args}}, type, acc, fun) do
     [struct, in_args, expr, callback] = args
-    {in_args, acc} = map_block_args(in_args, acc, fun)
+    {in_args, acc} = Enum.map_reduce(in_args, acc, fun)
 
     {expr, acc} =
       case type do
@@ -254,15 +254,4 @@ defmodule Nx.Defn.Tree do
       arg, acc -> {arg, acc}
     end)
   end
-
-  @doc false
-  def map_block_args(list, acc, fun) when is_list(list) do
-    Enum.map_reduce(list, acc, fn
-      %T{} = arg, acc ->
-        fun.(arg, acc)
-
-      arg, acc ->
-        {arg, acc}
-    end)
-  end
 end