Handle Nx.Block structs in LazyContainer/Composite for defn/JIT; trim LU block + EXLA subfn names

Author: Chapaman

exla/test/exla/nx_linalg_doctest_test.exs

@@ -25,7 +25,7 @@ defmodule EXLA.NxLinAlgDoctestTest do
     least_squares: 3,
     determinant: 1,
     matrix_power: 2,
-    lu: 2,
+    lu: 1,
     qr: 2
   ]
 

nx/lib/nx/block.ex

@@ -73,3 +73,35 @@ end
 defmodule Nx.Block.IFFT2 do
   defstruct eps: nil, lengths: nil, axes: nil
 end
+
+defmodule Nx.Block do
+  @moduledoc false
+
+  @structs [
+    Nx.Block.LogicalNot,
+    Nx.Block.Phase,
+    Nx.Block.Cholesky,
+    Nx.Block.Solve,
+    Nx.Block.QR,
+    Nx.Block.Eigh,
+    Nx.Block.SVD,
+    Nx.Block.LU,
+    Nx.Block.Determinant,
+    Nx.Block.AllClose,
+    Nx.Block.CumulativeSum,
+    Nx.Block.CumulativeProduct,
+    Nx.Block.CumulativeMin,
+    Nx.Block.CumulativeMax,
+    Nx.Block.Take,
+    Nx.Block.TakeAlongAxis,
+    Nx.Block.TopK,
+    Nx.Block.FFT2,
+    Nx.Block.IFFT2
+  ]
+
+  for mod <- @structs do
+    def block?(%unquote(mod){}), do: true
+  end
+
+  def block?(_), do: false
+end

nx/lib/nx/defn/composite.ex

@@ -36,7 +36,32 @@ defmodule Nx.Defn.Composite do
     |> Enum.all?(fn {l, r} -> compatible?(l, r, fun) end)
   end
 
-  def compatible?(%mod{} = left, %mod{} = right, fun) do
+  def compatible?(left, right, fun) when is_struct(left) and is_struct(right) do
+    cond do
+      Nx.Block.block?(left) and Nx.Block.block?(right) ->
+        left == right
+
+      left.__struct__ == right.__struct__ ->
+        compatible_struct(left, right, fun)
+
+      true ->
+        false
+    end
+  end
+
+  def compatible?(left, right, fun) when map_size(left) == map_size(right) do
+    Enum.all?(left, fn {k, v1} ->
+      case right do
+        %{^k => v2} -> compatible?(v1, v2, fun)
+        %{} -> false
+      end
+    end)
+  end
+
+  def compatible?(_, _, _),
+    do: false
+
+  defp compatible_struct(left, right, fun) do
     # LazyContainer is fully recursive but we don't want to go full recursive
     # unless we have to, so we can also compare structures along the way.
     {left, right} =
@@ -59,21 +84,6 @@ defmodule Nx.Defn.Composite do
     Enum.zip(left, right) |> Enum.all?(fn {l, r} -> compatible?(l, r, fun) end)
   end
 
-  def compatible?(%_{}, %_{}, _fun),
-    do: false
-
-  def compatible?(left, right, fun) when map_size(left) == map_size(right) do
-    Enum.all?(left, fn {k, v1} ->
-      case right do
-        %{^k => v2} -> compatible?(v1, v2, fun)
-        %{} -> false
-      end
-    end)
-  end
-
-  def compatible?(_, _, _),
-    do: false
-
   @doc """
   Counts the number of non-composite types in the composite type.
 
@@ -89,7 +99,14 @@ defmodule Nx.Defn.Composite do
   """
   def count(tree), do: count(tree, 0)
   defp count(tensor, acc) when is_tensor(tensor), do: acc + 1
-  defp count(container, acc), do: Nx.Container.reduce(container, acc, &count/2)
+
+  defp count(other, acc) do
+    if Nx.Block.block?(other) do
+      acc
+    else
+      Nx.Container.reduce(other, acc, &count/2)
+    end
+  end
 
   @doc """
   Traverses recursively the given composite types with `fun`.
@@ -117,8 +134,13 @@ defmodule Nx.Defn.Composite do
   def traverse(expr, acc, fun) when is_tensor(expr) and is_function(fun, 2),
     do: fun.(expr, acc)
 
-  def traverse(container, acc, fun),
-    do: Nx.Container.traverse(container, acc, &traverse(&1, &2, fun))
+  def traverse(expr, acc, fun) when is_function(fun, 2) do
+    if Nx.Block.block?(expr) do
+      {expr, acc}
+    else
+      Nx.Container.traverse(expr, acc, &traverse(&1, &2, fun))
+    end
+  end
 
   @doc """
   Reduces recursively the given composite types with `acc` and `fun`.
@@ -132,8 +154,13 @@ defmodule Nx.Defn.Composite do
   def reduce(expr, acc, fun) when is_tensor(expr) and is_function(fun, 2),
     do: fun.(expr, acc)
 
-  def reduce(container, acc, fun),
-    do: Nx.Container.reduce(container, acc, &reduce(&1, &2, fun))
+  def reduce(expr, acc, fun) when is_function(fun, 2) do
+    if Nx.Block.block?(expr) do
+      acc
+    else
+      Nx.Container.reduce(expr, acc, &reduce(&1, &2, fun))
+    end
+  end
 
   @doc """
   Flattens recursively the given list of composite types.
@@ -163,6 +190,13 @@ defmodule Nx.Defn.Composite do
        when is_number(number) or is_struct(number, Complex),
        do: [number | acc]
 
-  defp flatten_each(container, acc),
-    do: Nx.Container.reduce(container, acc, &flatten_each/2)
+  defp flatten_each(other, acc) do
+    cond do
+      Nx.Block.block?(other) ->
+        [other | acc]
+
+      true ->
+        Nx.Container.reduce(other, acc, &flatten_each/2)
+    end
+  end
 end

nx/lib/nx/lazy_container.ex

@@ -107,16 +107,20 @@ end
 
 defimpl Nx.LazyContainer, for: Any do
   def traverse(data, acc, fun) do
-    case Nx.Container.impl_for(data) do
-      nil ->
-        raise Protocol.UndefinedError,
-          protocol: @protocol,
-          value: data,
-          description:
-            "data-structures given to defn/Nx must implement either Nx.LazyContainer or Nx.Container"
-
-      impl ->
-        impl.traverse(data, acc, &Nx.LazyContainer.traverse(&1, &2, fun))
+    if Nx.Block.block?(data) do
+      {data, acc}
+    else
+      case Nx.Container.impl_for(data) do
+        nil ->
+          raise Protocol.UndefinedError,
+            protocol: @protocol,
+            value: data,
+            description:
+              "data-structures given to defn/Nx must implement either Nx.LazyContainer or Nx.Container"
+
+        impl ->
+          impl.traverse(data, acc, &Nx.LazyContainer.traverse(&1, &2, fun))
+      end
     end
   end
 end