refactor: only return the chain in the public interface

Author: polvalente

nx/lib/nx/defn/graph_splitter.ex

@@ -9,6 +9,12 @@ defmodule Nx.Defn.GraphSplitter do
   Traverses the expression and splits it into stages.
   """
   def traverse(expr, expr_split_fn \\ fn _ -> false end) do
+    {chain, _, _} = traverse_and_return_cache(expr, expr_split_fn)
+    chain
+  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}
     # that we can then compile and chain-execute elsewhere. category is either :gather, :reduce or :none
     state = %{

nx/test/nx/defn/sharding_compiler/passes/graph_splitter_test.exs

@@ -23,7 +23,7 @@ defmodule Nx.Defn.GraphSplitterTest do
         _ -> false
       end
 
-      {chain, cache, state} = GraphSplitter.traverse(expr, split_fn)
+      {chain, cache, state} = GraphSplitter.traverse_and_return_cache(expr, split_fn)
 
       assert [
                %Stage{
@@ -143,7 +143,7 @@ defmodule Nx.Defn.GraphSplitterTest do
         _ -> false
       end
 
-      {chain, cache, state} = GraphSplitter.traverse(expr, split_fn)
+      {chain, cache, state} = GraphSplitter.traverse_and_return_cache(expr, split_fn)
 
       assert [
                %Stage{
@@ -292,8 +292,7 @@ defmodule Nx.Defn.GraphSplitterTest do
         _ -> false
       end
 
-      assert {[%Stage{} = stage_0, %Stage{} = stage_1], _cache, _state} =
-               GraphSplitter.traverse(expr, split_fn)
+      assert [%Stage{} = stage_0, %Stage{} = stage_1] = GraphSplitter.traverse(expr, split_fn)
 
       [{arg1_id, %T{shape: {2, 3}, type: {:u, 8}, data: %Expr{args: [0]}}}] =
         Enum.to_list(stage_0.arguments)
@@ -358,8 +357,7 @@ defmodule Nx.Defn.GraphSplitterTest do
         _ -> false
       end
 
-      assert {[%Stage{} = stage_0, %Stage{} = stage_1], _cache, _state} =
-               GraphSplitter.traverse(expr, split_fn)
+      assert [%Stage{} = stage_0, %Stage{} = stage_1] = GraphSplitter.traverse(expr, split_fn)
 
       [{arg1_id, %T{shape: {2, 3}, type: {:u, 8}, data: %Expr{args: [0]}}}] =
         Enum.to_list(stage_0.arguments)