Add least_squares function (#1438)

elixir-nx · Jan 30, 2024 · 0b97ba8 · 0b97ba8
1 parent 3e42ec3
commit 0b97ba8
Show file tree

Hide file tree

Showing 4 changed files with 139 additions and 2 deletions.
diff --git a/exla/test/exla/mlir/nx_linalg_doctest_test.exs b/exla/test/exla/mlir/nx_linalg_doctest_test.exs
@@ -15,7 +15,7 @@ defmodule EXLA.MLIR.NxLinAlgDoctestTest do
     invert: 1,
     matrix_power: 2
   ]
-  @rounding_error_doctests [triangular_solve: 3, eigh: 2, cholesky: 1]
+  @rounding_error_doctests [triangular_solve: 3, eigh: 2, cholesky: 1, least_squares: 2]
 
   @excluded_doctests @function_clause_error_doctests ++
                        @rounding_error_doctests ++

diff --git a/nx/lib/nx/lin_alg.ex b/nx/lib/nx/lin_alg.ex
@@ -2138,6 +2138,105 @@ defmodule Nx.LinAlg do
     Nx.sum(s > tol)
   end
 
+  @doc """
+  Return the least-squares solution to a linear matrix equation Ax = b.
+
+  ## Examples
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([[1, 2], [2, 3]]), Nx.tensor([1, 2]))
+      #Nx.Tensor<
+        f32[2]
+        [1.0000004768371582, -2.665601925855299e-7]
+      >
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([[0, 1], [1, 1], [2, 1], [3, 1]]), Nx.tensor([-1, 0.2, 0.9, 2.1]))
+      #Nx.Tensor<
+        f32[2]
+        [0.9966150522232056, -0.9479667544364929]
+      >
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([[1, 2, 3], [4, 5, 6]]), Nx.tensor([1, 2]))
+      #Nx.Tensor<
+        f32[3]
+        [-0.05531728267669678, 0.11113593727350235, 0.27758902311325073]
+      >
+
+  ## Error cases
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([1, 2, 3]), Nx.tensor([1, 2]))
+      ** (ArgumentError) tensor of 1st argument must have rank 2, got rank 1 with shape {3}
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([[1, 2], [2, 3]]), Nx.tensor([[1, 2], [3, 4]]))
+      ** (ArgumentError) tensor of 2nd argument must have rank 1, got rank 2 with shape {2, 2}
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([[1, Complex.new(0, 2)], [3, Complex.new(0, -4)]]),  Nx.tensor([1, 2]))
+      ** (ArgumentError) Nx.LinAlg.least_squares/2 is not yet implemented for complex inputs
+
+      iex> Nx.LinAlg.least_squares(Nx.tensor([[1, 2], [2, 3]]), Nx.tensor([1, 2, 3]))
+      ** (ArgumentError) the number of rows of the matrix as the 1st argument and the number of columns of the vector as the 2nd argument must be the same, got 1st argument shape {2, 2} and 2nd argument shape {3}
+  """
+  @doc from_backend: false
+  defn least_squares(a, b) do
+    %T{type: a_type, shape: a_shape} = Nx.to_tensor(a)
+    a_size = Nx.rank(a_shape)
+    %T{type: b_type, shape: b_shape} = Nx.to_tensor(b)
+    b_size = Nx.rank(b_shape)
+
+    case a_type do
+      {:c, _} ->
+        raise ArgumentError, "Nx.LinAlg.least_squares/2 is not yet implemented for complex inputs"
+
+      _ ->
+        nil
+    end
+
+    case b_type do
+      {:c, _} ->
+        raise ArgumentError, "Nx.LinAlg.least_squares/2 is not yet implemented for complex inputs"
+
+      _ ->
+        nil
+    end
+
+    if a_size != 2 do
+      raise(
+        ArgumentError,
+        "tensor of 1st argument must have rank 2, got rank #{inspect(a_size)} with shape #{inspect(a_shape)}"
+      )
+    end
+
+    if b_size != 1 do
+      raise(
+        ArgumentError,
+        "tensor of 2nd argument must have rank 1, got rank #{inspect(b_size)} with shape #{inspect(b_shape)}"
+      )
+    end
+
+    {a1, _a2} = a_shape
+    {b1} = b_shape
+
+    if a1 != b1 do
+      raise(
+        ArgumentError,
+        "the number of rows of the matrix as the 1st argument and " <>
+          "the number of columns of the vector as the 2nd argument must be the same, " <>
+          "got 1st argument shape #{inspect(a_shape)} and 2nd argument shape #{inspect(b_shape)}"
+      )
+    end
+
+    case a_shape do
+      {m, n} when m == n ->
+        Nx.LinAlg.solve(a, b)
+
+      {m, n} when m != n ->
+        Nx.LinAlg.pinv(a, eps: 1.0e-15)
+        |> Nx.dot(b)
+
+      _ ->
+        nil
+    end
+  end
+
   defp apply_vectorized(tensor, fun) when is_function(fun, 1) do
     # same as Nx's apply_vectorized defp, but written in a "public-api" way!
     %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor)

diff --git a/nx/test/nx/lin_alg_test.exs b/nx/test/nx/lin_alg_test.exs
@@ -975,4 +975,41 @@ defmodule Nx.LinAlgTest do
       end
     end)
   end
+
+  describe "least_squares" do
+    test "properties for linear equations" do
+      key = Nx.Random.key(System.unique_integer())
+
+      # Calucate linear equations Ax = y by using least-squares solution
+      for {m, n} <- [{2, 2}, {3, 2}, {4, 3}], reduce: key do
+        key ->
+          # Generate x as temporary solution and A as base matrix
+          {a_base, key} = Nx.Random.randint(key, 1, 10, shape: {m, n})
+          {x_temp, key} = Nx.Random.randint(key, 1, 10, shape: {n})
+
+          # Generate y as base vector by x and A
+          # to prepare an equation that can be solved exactly
+          y_base = Nx.dot(a_base, x_temp)
+
+          # Generate y as random noise vector and A as random noise matrix
+          noise_eps = 1.0e-2
+          {a_noise, key} = Nx.Random.uniform(key, 0, noise_eps, shape: {m, n})
+          {y_noise, key} = Nx.Random.uniform(key, 0, noise_eps, shape: {m})
+
+          # Add noise to prepare equations that cannot be solved without approximation,
+          # such as the least-squares method
+          a = Nx.add(a_base, a_noise)
+          y = Nx.add(y_base, y_noise)
+
+          # Calculate least-squares solution to a linear matrix equation Ax = y
+          x = Nx.LinAlg.least_squares(a, y)
+
+          # Check linear matrix equation
+          Nx.dot(a, x)
+          |> assert_all_close(y, atol: noise_eps * 10)
+
+          key
+      end
+    end
+  end
 end
diff --git a/torchx/test/torchx/nx_linalg_doctest_test.exs b/torchx/test/torchx/nx_linalg_doctest_test.exs
@@ -17,7 +17,8 @@ defmodule Torchx.NxLinAlgDoctestTest do
     solve: 2,
     invert: 1,
     determinant: 1,
-    pinv: 2
+    pinv: 2,
+    least_squares: 2
   ]
 
   # Results do not match but properties are respected