add split function

nx/.gitignore

@@ -21,3 +21,6 @@ erl_crash.dump
 
 # Ignore package tarball (built via "mix hex.build").
 nx-*.tar
+
+# ASDF files
+.tool-versions

nx/lib/nx.ex

@@ -3265,6 +3265,124 @@ defmodule Nx do
     end)
   end
 
+  @doc ~S"""
+  Split a tensor into train and test subsets.
+
+  `split` split must be either an integer greater than zero and less than the length of the tensor or a float number between `0.0` and `1.0`.
+
+  ## Options
+
+    * `:axis` - The axis along which to take the values from. Defaults to `0`.
+
+  ## Examples
+
+    Split a tensor into two separate tensors.
+
+    iex> {train, test} = Nx.split(Nx.tensor([[3, 6, 5], [26, 75, 3], [23, 4, 1]]), 2, axis: 0)
+    iex> train
+    #Nx.Tensor<
+      s64[2][3]
+      [
+        [3, 6, 5],
+        [26, 75, 3]
+      ]
+    >
+    iex> test
+    #Nx.Tensor<
+      s64[1][3]
+      [
+        [23, 4, 1]
+      ]
+    >
+
+    iex> {train, test} = Nx.split(Nx.tensor([[3, 6, 5], [26, 75, 3], [23, 4, 1]]), 2, axis: 1)
+    iex> train
+    #Nx.Tensor<
+      s64[3][2]
+      [
+        [3, 6],
+        [26, 75],
+        [23, 4]
+      ]
+    >
+    iex> test
+    #Nx.Tensor<
+      s64[3][1]
+      [
+        [5],
+        [3],
+        [1]
+      ]
+    >
+
+    iex> {train, test} = Nx.split(Nx.tensor([[3, 6, 5, 9], [26, 75, 3, 20], [23, 4, 1, 56], [40, 6, 78, 94]]), 0.5, axis: 0)
+    iex> train
+    #Nx.Tensor<
+      s64[2][4]
+      [
+        [3, 6, 5, 9],
+        [26, 75, 3, 20]
+      ]
+    >
+    iex> test
+    #Nx.Tensor<
+      s64[2][4]
+      [
+        [23, 4, 1, 56],
+        [40, 6, 78, 94]
+      ]
+    >
+  """
+  @doc type: :indexed
+  def split(tensor, split, opts \\ [])
+
+  def split(%T{shape: shape} = tensor, split, opts) when is_integer(split) do
+    opts = keyword!(opts, axis: 0)
+    axis = Keyword.fetch!(opts, :axis)
+
+    if split > 0 do
+      axis = find_axis(tensor, axis)
+      values = elem(shape, axis)
+      size = values - split
+
+      {
+        slice_along_axis(tensor, 0, split, axis: axis),
+        slice_along_axis(tensor, split, size, axis: axis)
+      }
+    else
+      raise "split must be an integer greater than zero and less than the length of the tensor."
+    end
+  end
+
+  def split(%T{shape: shape} = tensor, split, opts) when is_float(split) do
+    opts = keyword!(opts, axis: 0)
+    axis = Keyword.fetch!(opts, :axis)
+
+    if split > 0.0 and split < 1.0 do
+      axis = find_axis(tensor, axis)
+
+      values = elem(shape, axis)
+
+      split_size = Kernel.ceil(split * values)
+
+      split_size =
+        cond do
+          split_size < 1 -> 1
+          split_size >= values -> 1
+          true -> split_size
+        end
+
+      remaining_size = values - split_size
+
+      {
+        slice_along_axis(tensor, 0, split_size, axis: axis),
+        slice_along_axis(tensor, split_size, remaining_size, axis: axis)
+      }
+    else
+      raise "split must be a float number between 0.0 and 1.0."
+    end
+  end
+
   @doc """
   Broadcasts `tensor` to the given `broadcast_shape`.
 
@@ -16061,4 +16179,21 @@ defmodule Nx do
       end)
     end
   end
+
+  defp find_axis(tensor, axis) do
+    axis_values = axes(tensor)
+    axis_names = names(tensor)
+
+    cond do
+      is_integer(axis) and axis in axis_values ->
+        axis
+
+      is_atom(axis) and axis in axis_names ->
+        dimensions = Enum.zip(axis_names, axis_values)
+        dimensions[axis]
+
+      true ->
+        raise ":axis is out of tensor bounds."
+    end
+  end
 end

nx/test/nx_test.exs

@@ -3022,4 +3022,105 @@ defmodule NxTest do
       assert_all_close(zeros, Nx.imag(x_ifft), atol: 1.0e-8)
     end
   end
+
+  describe "split/2" do
+    test "split is less than zero" do
+      tensor = Nx.iota({10, 2}, names: [:x, :y])
+
+      assert_raise RuntimeError,
+                   "split must be an integer greater than zero and less than the length of the tensor.",
+                   fn ->
+                     Nx.split(tensor, -1)
+                   end
+    end
+
+    test "split is a float out of bounds" do
+      tensor = Nx.iota({10, 2}, names: [:x, :y])
+
+      assert_raise RuntimeError,
+                   "split must be a float number between 0.0 and 1.0.",
+                   fn ->
+                     Nx.split(tensor, 1.0)
+                   end
+    end
+
+    test "split is greater than tensor length" do
+      tensor = Nx.iota({10, 2}, names: [:x, :y])
+
+      assert_raise ArgumentError,
+                   "length at axis 1 must be less than axis size of 2, got: 3",
+                   fn ->
+                     Nx.split(tensor, 3, axis: 1)
+                   end
+    end
+
+    test "axis is out of tensor bounds" do
+      tensor = Nx.iota({10, 2}, names: [:x, :y])
+
+      assert_raise RuntimeError,
+                   ":axis is out of tensor bounds.",
+                   fn ->
+                     Nx.split(tensor, 2, axis: 2)
+                   end
+    end
+
+    test "named axis is out of tensor bounds" do
+      tensor = Nx.iota({10, 2}, names: [:x, :y])
+
+      assert_raise RuntimeError,
+                   ":axis is out of tensor bounds.",
+                   fn ->
+                     Nx.split(tensor, 2, axis: :z)
+                   end
+    end
+
+    test "split into 50% for training and 50% for testing with floats on columns" do
+      tensor = Nx.iota({4, 4}, names: [:rows, :columns])
+      {train, test} = Nx.split(tensor, 0.5, axis: :columns)
+
+      assert {4, 2} == Nx.shape(train)
+      assert {4, 2} == Nx.shape(test)
+    end
+
+    test "split into 70% for training and 30% for testing along a named :axis" do
+      tensor = Nx.iota({100, 6}, names: [:rows, :columns])
+      {train, test} = Nx.split(tensor, 70, axis: :rows)
+
+      assert {70, 6} == Nx.shape(train)
+      assert {30, 6} == Nx.shape(test)
+    end
+
+    test "split into 90% for training and 10% for testing along a named :axis" do
+      tensor = Nx.iota({2, 100}, names: [:rows, :columns])
+      {train, test} = Nx.split(tensor, 90, axis: :columns)
+
+      assert {2, 90} == Nx.shape(train)
+      assert {2, 10} == Nx.shape(test)
+    end
+
+    test "split into 50% for training and 50% for testing along the :axis 1" do
+      tensor = Nx.iota({100, 10})
+      {train, test} = Nx.split(tensor, 5, axis: 1)
+
+      assert {100, 5} == Nx.shape(train)
+      assert {100, 5} == Nx.shape(test)
+    end
+
+    test "split into 61% for training and 39% for testing" do
+      tensor = Nx.iota({100, 10})
+      {train, test} = Nx.split(tensor, 61)
+
+      assert {61, 10} == Nx.shape(train)
+      assert {39, 10} == Nx.shape(test)
+    end
+
+    test "split into 60% for training and 40% for testing with unbalanced data" do
+      tensor = Nx.iota({99, 4})
+
+      {train, test} = Nx.split(tensor, 60)
+
+      assert {60, 4} == Nx.shape(train)
+      assert {39, 4} == Nx.shape(test)
+    end
+  end
 end