feat: add GraphSplitter.run/2

Author: polvalente

nx/lib/nx/defn/graph_splitter.ex

@@ -13,6 +13,58 @@ defmodule Nx.Defn.GraphSplitter do
     chain
   end
 
+  @doc """
+  Executes the stage chain with the given arguments.
+  """
+  def run(chain, args) do
+    scope =
+      Enum.with_index(args, fn arg, idx -> {{nil, idx}, arg} end)
+      |> Map.new()
+
+    scope =
+      Enum.reduce(chain, scope, fn stage, scope ->
+        %{id: id, expr: expr, argument_sources: argument_sources, arguments: arguments} = stage
+
+        args =
+          arguments
+          |> Enum.map(fn {id, %T{data: %Expr{args: [idx]}}} ->
+            source = Map.fetch!(argument_sources, id)
+            argument = Map.fetch!(scope, source)
+            {idx, argument}
+          end)
+          |> Enum.sort_by(fn {idx, _} -> idx end)
+          |> Enum.map(fn {_, argument} -> argument end)
+
+        case Nx.Defn.jit_apply(fn _ -> expr end, [List.to_tuple(args)]) do
+          %T{} = tensor ->
+            Map.put(scope, {id, 0}, tensor)
+
+          tuple ->
+            {_idx, scope} =
+              tuple
+              |> Tuple.to_list()
+              |> Enum.reduce({0, scope}, fn tensor, {idx, scope} ->
+                {idx + 1, Map.put(scope, {id, idx}, tensor)}
+              end)
+
+            scope
+        end
+      end)
+
+    last_stage = List.last(chain)
+
+    if is_tuple(last_stage.expr) do
+      scope
+      |> Enum.filter(fn {{id, _}, _} -> id == last_stage.id end)
+      |> Enum.sort_by(fn {{_, idx}, _} -> idx end)
+      |> Enum.map(fn {_, tensor} -> tensor end)
+      |> List.to_tuple()
+    else
+      {_, tensor} = Enum.find(scope, fn {{id, _}, _} -> id == last_stage.id end)
+      tensor
+    end
+  end
+
   @doc false
   def traverse_and_return_cache(expr, expr_split_fn) do
     # expression_chain is going to be a reverse-accumulation of {category, subexpr}

nx/test/nx/defn/graph_splitter_test.exs

@@ -395,4 +395,113 @@ defmodule Nx.Defn.GraphSplitterTest do
       assert %T{data: %Expr{id: ^arg_1_id, op: :parameter, args: [1]}} = a
     end
   end
+
+  describe "run/2" do
+    test "executes the stages chain and returns the correct result" do
+      function = fn arg0, arg1 ->
+        # root
+        x = Nx.multiply(arg0, arg1) |> Nx.Defn.Expr.metadata(%{split: true})
+
+        # left side
+        w_left = Nx.multiply(x, arg1) |> Nx.Defn.Expr.metadata(%{split: true})
+
+        # right side
+        w_right = Nx.divide(x, arg1) |> Nx.Defn.Expr.metadata(%{split: true})
+
+        # merge
+        Nx.add(w_right, w_left)
+      end
+
+      args = [Nx.tensor([1, 2]), Nx.tensor([3, 4])]
+
+      # This is used in the final assertion of this test
+      expected_result = Nx.Defn.jit_apply(function, args)
+
+      expr = apply(Nx.Defn.debug_expr(function), args)
+
+      split_fn = fn
+        %T{data: %Expr{op: :metadata, args: [_expr, %{split: true}]}} -> true
+        _ -> false
+      end
+
+      chain = GraphSplitter.traverse(expr, split_fn)
+
+      assert [root, side1, side2, merge] = chain
+
+      assert {%T{data: %Expr{op: :multiply, args: [arg0, arg1]}}} = root.expr
+      assert %T{data: %Expr{op: :parameter, args: [0]}} = arg0
+      assert %T{data: %Expr{op: :parameter, args: [1]}} = arg1
+
+      # because things are balanced, we don't know which of side1 and side2 are left and right
+      # in our expr, so we should disambiguate:
+
+      {[%Stage{} = left], [%Stage{} = right]} =
+        Enum.split_with([side1, side2], fn %Stage{expr: {expr}} -> expr.data.op == :multiply end)
+
+      # left should depend on exactly the same parameters as the root, as it's pulling from
+      # the global scope
+      assert {%T{data: %Expr{op: :multiply, args: [x, arg1_left]}}} = left.expr
+
+      assert %T{
+               data: %Expr{
+                 op: :metadata,
+                 args: [
+                   %T{data: %Expr{id: x_left_id, op: :parameter, args: [1]}},
+                   %{split: true}
+                 ]
+               }
+             } = x
+
+      assert %T{data: %Expr{id: arg1_left_id, op: :parameter, args: [0]}} = arg1_left
+
+      assert left.argument_sources[arg1_left_id] == {nil, 1}
+      assert left.argument_sources[x_left_id] == {root.id, 0}
+
+      # right should depend on the result of the root and on arg1, but arg1 will be reindexed
+      # we should assert that the argument source for arg1_right is correct
+      assert {%T{data: %Expr{op: :divide, args: [x, arg1_right]}}} = right.expr
+
+      assert %T{
+               data: %Expr{
+                 op: :metadata,
+                 args: [
+                   %T{data: %Expr{id: x_right_id, op: :parameter, args: [1]}},
+                   %{split: true}
+                 ]
+               }
+             } = x
+
+      assert %T{data: %Expr{id: arg1_right_id, op: :parameter, args: [0]}} = arg1_right
+
+      assert right.argument_sources[arg1_right_id] == {nil, 1}
+      assert right.argument_sources[x_right_id] == {root.id, 0}
+
+      assert %T{data: %Expr{op: :add, args: [w_right, w_left]}} = merge.expr
+
+      assert %T{
+               data: %Expr{
+                 op: :metadata,
+                 args: [
+                   %T{data: %Expr{id: w_right_id, op: :parameter, args: [0]}},
+                   %{split: true}
+                 ]
+               }
+             } = w_right
+
+      assert %T{
+               data: %Expr{
+                 op: :metadata,
+                 args: [
+                   %T{data: %Expr{id: w_left_id, op: :parameter, args: [1]}},
+                   %{split: true}
+                 ]
+               }
+             } = w_left
+
+      assert merge.argument_sources[w_right_id] == {right.id, 0}
+      assert merge.argument_sources[w_left_id] == {left.id, 0}
+
+      assert GraphSplitter.run(chain, args) == expected_result
+    end
+  end
 end