Remove Nx.Container and Block.name from block.ex, update Lu.lu from /2 to /1

Author: Chapaman

exla/lib/exla/defn.ex

@@ -763,20 +763,18 @@ defmodule EXLA.Defn do
 
   defp cached_recur_operator(:block, %T{data: %Expr{args: args}}, state, cache) do
     [struct, in_args, expr, _callback] = args
-    op = Nx.Block.name(struct)
 
-    {call_prefix, _opts_suffix} = Enum.split_while(in_args, &(not is_list(&1)))
-
-    {call_args, cache} = Enum.map_reduce(call_prefix, cache, &recur_operator(&1, state, &2))
-    key = computation_key(op, [struct | call_args])
+    {call_args, cache} = Enum.map_reduce(in_args, cache, &recur_operator(&1, state, &2))
+    key = computation_key(struct.__struct__, [struct | call_args])
 
     {call_body, cache} =
       case cache do
         %{^key => computation} ->
           {computation, cache}
 
         %{} ->
-          {computation, cache} = block_computation("block", call_args, expr, state, cache)
+          {computation, cache} =
+            block_computation(block_subfunction_description(struct), call_args, expr, state, cache)
           {computation, Map.put(cache, key, computation)}
       end
 
@@ -1831,8 +1829,15 @@ defmodule EXLA.Defn do
     {region, merge_outfeed(cache, comp_cache)}
   end
 
-  defp block_computation(name, args, expr, %{builder: %Function{}} = state, cache) do
-    %Function{module: module, name: name} = subbuilder(state.builder, name)
+  defp block_subfunction_description(%_{} = struct) do
+    struct.__struct__
+    |> Module.split()
+    |> List.last()
+    |> Macro.underscore()
+  end
+
+  defp block_computation(description, args, expr, %{builder: %Function{}} = state, cache) do
+    %Function{module: module, name: name} = subbuilder(state.builder, description)
 
     arg_typespecs = Enum.map(args, &Value.get_typespec/1)
     out_typespecs = container_to_typespecs(expr)

nx/lib/nx/block.ex

@@ -15,103 +15,61 @@ defmodule Nx.Block.Solve do
 end
 
 defmodule Nx.Block.QR do
-  @derive {Nx.Container, containers: [], keep: [:eps, :mode]}
   defstruct eps: 1.0e-10, mode: :reduced
 end
 
 defmodule Nx.Block.Eigh do
-  @derive {Nx.Container, containers: [], keep: [:max_iter, :eps]}
   defstruct max_iter: 1000, eps: 1.0e-4
 end
 
 defmodule Nx.Block.SVD do
-  @derive {Nx.Container, containers: [], keep: [:max_iter, :full_matrices?]}
   defstruct max_iter: 100, full_matrices?: true
 end
 
 defmodule Nx.Block.LU do
-  @derive {Nx.Container, containers: [], keep: [:eps]}
-  defstruct eps: 1.0e-10
+  defstruct []
 end
 
 defmodule Nx.Block.Determinant do
   defstruct []
 end
 
 defmodule Nx.Block.AllClose do
-  @derive {Nx.Container, containers: [], keep: [:equal_nan, :rtol, :atol]}
   defstruct equal_nan: false, rtol: 1.0e-5, atol: 1.0e-8
 end
 
 defmodule Nx.Block.CumulativeSum do
-  @derive {Nx.Container, containers: [], keep: [:axis, :reverse]}
   defstruct axis: 0, reverse: false
 end
 
 defmodule Nx.Block.CumulativeProduct do
-  @derive {Nx.Container, containers: [], keep: [:axis, :reverse]}
   defstruct axis: 0, reverse: false
 end
 
 defmodule Nx.Block.CumulativeMin do
-  @derive {Nx.Container, containers: [], keep: [:axis, :reverse]}
   defstruct axis: 0, reverse: false
 end
 
 defmodule Nx.Block.CumulativeMax do
-  @derive {Nx.Container, containers: [], keep: [:axis, :reverse]}
   defstruct axis: 0, reverse: false
 end
 
 defmodule Nx.Block.Take do
-  @derive {Nx.Container, containers: [], keep: [:axis]}
   defstruct axis: 0
 end
 
 defmodule Nx.Block.TakeAlongAxis do
-  @derive {Nx.Container, containers: [], keep: [:axis]}
   defstruct axis: 0
 end
 
 defmodule Nx.Block.TopK do
-  @derive {Nx.Container, containers: [], keep: [:k]}
   defstruct k: 1
 end
 
 defmodule Nx.Block.FFT2 do
-  @derive {Nx.Container, containers: [], keep: [:eps, :lengths, :axes]}
   defstruct eps: nil, lengths: nil, axes: nil
 end
 
 defmodule Nx.Block.IFFT2 do
-  @derive {Nx.Container, containers: [], keep: [:eps, :lengths, :axes]}
   defstruct eps: nil, lengths: nil, axes: nil
 end
-
-defmodule Nx.Block do
-  @moduledoc false
-
-  def name(%{__struct__: module}) do
-    case module do
-      Nx.Block.LogicalNot -> :logical_not
-      Nx.Block.Phase -> :phase
-      Nx.Block.AllClose -> :all_close
-      Nx.Block.CumulativeSum -> :cumulative_sum
-      Nx.Block.CumulativeProduct -> :cumulative_product
-      Nx.Block.CumulativeMin -> :cumulative_min
-      Nx.Block.CumulativeMax -> :cumulative_max
-      Nx.Block.Cholesky -> :cholesky
-      Nx.Block.Solve -> :solve
-      Nx.Block.QR -> :qr
-      Nx.Block.Eigh -> :eigh
-      Nx.Block.SVD -> :svd
-      Nx.Block.LU -> :lu
-      Nx.Block.Determinant -> :determinant
-      Nx.Block.Take -> :take
-      Nx.Block.TakeAlongAxis -> :take_along_axis
-      Nx.Block.TopK -> :top_k
-      Nx.Block.FFT2 -> :fft2
-      Nx.Block.IFFT2 -> :ifft2
-    end
-  end
-end

nx/lib/nx/defn/composite.ex

@@ -3,7 +3,7 @@ defmodule Nx.Defn.Composite do
   Functions to deal with composite data types.
 
   Composite data-types are traversed according to `Nx.Container`.
-  If a regular tensor is given, it is individually traversed. 
+  If a regular tensor is given, it is individually traversed.
   Numerical values, such as integers, floats, and complex numbers
   are not normalized before hand. Use `Nx.to_tensor/1` to do so.
 

nx/lib/nx/defn/evaluator.ex

@@ -136,10 +136,8 @@ defmodule Nx.Defn.Evaluator do
   defp compute_cache(:block, %{data: %Expr{args: args}}, state, cache) do
     [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(Nx.Block.name(struct), call_prefix)
+    {in_args, cache} = Enum.map_reduce(in_args, cache, &compute_cache(&1, state, &2))
+    key = computation_key(struct.__struct__, in_args)
 
     {{expr, expr_cache}, cache} =
       case cache do
@@ -366,17 +364,16 @@ defmodule Nx.Defn.Evaluator do
 
   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))
-    {param_prefix, _} = Enum.split_while(in_args, &(not is_list(&1)))
-    backend = Nx.Shared.list_impl!(param_prefix)
+    backend = Nx.Shared.list_impl!(in_args)
 
     out =
       case ans do
         %{type: {:tuple, _}} -> expr
         _ -> ans
       end
 
-    fun = block_default_fun(expr, state, expr_cache, length(param_prefix))
-    {backend.block(struct, out, param_prefix, fun), caches}
+    fun = block_default_fun(expr, state, expr_cache, length(in_args))
+    {backend.block(struct, out, in_args, fun), caches}
   end
 
   defp eval_apply(:runtime_call, [expr, fun, out_template, opts], _ans, state, caches) do

nx/lib/nx/defn/expr.ex

@@ -44,9 +44,9 @@ defmodule Nx.Defn.Expr do
     * `runtime_call(out, tensor_or_container, opts, fun)`
 
     * `block(struct, block_args, default_expr, fun)` - `struct` is an `Nx.Block.*`
-      value, `block_args` are the tensors and keyword options passed to `Nx.block/4`,
-      `default_expr` is the traced default implementation, and `fun` is the block
-      callback
+      value, `block_args` are tensor expressions passed to `Nx.block/4` (options
+      live on `struct`), `default_expr` is the traced default implementation, and
+      `fun` is the block callback
 
   `defn` compilers must handle said nodes accordingly.
   """
@@ -426,10 +426,9 @@ defmodule Nx.Defn.Expr do
   end
 
   defp expr_block(struct, in_args, fun) do
-    {args, opts} = Enum.split_while(in_args, &(not is_list(&1)))
-    params = Enum.with_index(args, &parameter/2)
+    params = Enum.with_index(in_args, &parameter/2)
 
-    case apply(fun, [struct | params ++ opts]) do
+    case apply(fun, [struct | params]) do
       %{data: %{context: context}} = res ->
         expr(res, context, :block, [struct, in_args, res, fun])
 

nx/lib/nx/lin_alg.ex

@@ -1527,10 +1527,6 @@ defmodule Nx.LinAlg do
   @doc """
   Calculates the A = PLU decomposition of batched square 2-D matrices A.
 
-  ## Options
-
-    * `:eps` - Rounding error threshold that can be applied during the factorization
-
   ## Examples
 
       iex> {p, l, u} = Nx.LinAlg.lu(Nx.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
@@ -1734,8 +1730,7 @@ defmodule Nx.LinAlg do
       iex> Nx.LinAlg.lu(Nx.tensor([[1, 1, 1, 1], [-1, 4, 4, -1], [4, -2, 2, 0]]))
       ** (ArgumentError) tensor must be a square matrix or a batch of square matrices, got shape: {3, 4}
   """
-  def lu(tensor, opts \\ []) do
-    opts = keyword!(opts, eps: 1.0e-10)
+  def lu(tensor) do
     %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor)
     %T{type: type, shape: shape} = tensor = Nx.devectorize(tensor)
 
@@ -1748,8 +1743,8 @@ defmodule Nx.LinAlg do
        %{tensor | type: output_type, shape: l_shape, names: names},
        %{tensor | type: output_type, shape: u_shape, names: names}}
 
-    Nx.block(struct(Nx.Block.LU, opts), [tensor], output, fn %Nx.Block.LU{} = opts, t ->
-      Nx.LinAlg.LU.lu(opts, t)
+    Nx.block(%Nx.Block.LU{}, [tensor], output, fn %Nx.Block.LU{}, t ->
+      Nx.LinAlg.LU.lu(t)
     end)
     |> Nx.vectorize(vectorized_axes)
   end

nx/lib/nx/lin_alg/lu.ex

@@ -2,7 +2,7 @@ defmodule Nx.LinAlg.LU do
   @moduledoc false
   import Nx.Defn
 
-  defn lu(%Nx.Block.LU{} = _opts, a) do
+  defn lu(a) do
     vectorized_axes = a.vectorized_axes
 
     result =

torchx/lib/torchx/backend.ex

@@ -74,7 +74,7 @@ defmodule Torchx.Backend do
           qr_impl(hd(args), mode: struct.mode, eps: struct.eps)
 
         Nx.Block.LU ->
-          lu_impl(hd(args), eps: struct.eps)
+          lu_impl(hd(args))
 
         Nx.Block.Eigh ->
           eigh_impl(hd(args), max_iter: struct.max_iter, eps: struct.eps)
@@ -1224,7 +1224,7 @@ defmodule Torchx.Backend do
     end
   end
 
-  defp lu_impl(tensor, _opts) do
+  defp lu_impl(tensor) do
     tensor =
       if Nx.Type.integer?(tensor.type) do
         Nx.as_type(tensor, {:f, 32})