Fix(exla): triangular_solve with batched matrix input

exla/lib/exla/defn.ex

@@ -782,26 +782,27 @@ defmodule EXLA.Defn do
     lower = Keyword.fetch!(opts, :lower)
     transform = Keyword.fetch!(opts, :transform_a)
 
-    case Value.get_typespec(b).shape do
-      {dim} ->
-        b_shape = {dim, 1}
+    a_shape = Value.get_typespec(a).shape
+    b_shape = Value.get_typespec(b).shape
 
-        b =
-          b
-          |> to_type(type)
-          |> Value.reshape(Typespec.tensor(type, b_shape))
+    if tuple_size(a_shape) > tuple_size(b_shape) do
+      b_shape = Tuple.insert_at(b_shape, tuple_size(b_shape), 1)
 
-        typespec = Typespec.tensor(type, b_shape)
+      b =
+        b
+        |> to_type(type)
+        |> Value.reshape(Typespec.tensor(type, b_shape))
 
-        to_type(a, type)
-        |> Value.triangular_solve(b, left_side, lower, transform, typespec)
-        |> Value.reshape(Typespec.tensor(type, ans.shape))
+      typespec = Typespec.tensor(type, b_shape)
 
-      _ ->
-        typespec = Typespec.tensor(type, ans.shape)
+      to_type(a, type)
+      |> Value.triangular_solve(b, left_side, lower, transform, typespec)
+      |> Value.reshape(Typespec.tensor(type, ans.shape))
+    else
+      typespec = Typespec.tensor(type, ans.shape)
 
-        to_type(a, type)
-        |> Value.triangular_solve(to_type(b, type), left_side, lower, transform, typespec)
+      to_type(a, type)
+      |> Value.triangular_solve(to_type(b, type), left_side, lower, transform, typespec)
     end
   end
 

exla/test/exla/nx_linalg_doctest_test.exs

@@ -2,6 +2,11 @@ defmodule EXLA.NxLinAlgDoctestTest do
   use EXLA.Case, async: true
   import Nx, only: :sigils
 
+  setup do
+    Nx.default_backend(EXLA.Backend)
+    :ok
+  end
+
   @function_clause_error_doctests [
     solve: 2,
     triangular_solve: 3
@@ -15,7 +20,8 @@ defmodule EXLA.NxLinAlgDoctestTest do
     least_squares: 3,
     determinant: 1,
     matrix_power: 2,
-    lu: 2
+    lu: 2,
+    qr: 2
   ]
 
   @excluded_doctests @function_clause_error_doctests ++
@@ -402,4 +408,93 @@ defmodule EXLA.NxLinAlgDoctestTest do
       end
     end
   end
+
+  describe "triangular_solve" do
+    test "works with batched input" do
+      a =
+        Nx.tensor([
+          [
+            [-1, 0, 0],
+            [1, 1, 0],
+            [1, 1, 1]
+          ],
+          [
+            [2, 0, 0],
+            [4, -2, 0],
+            [-5, 1, 3]
+          ]
+        ])
+
+      b =
+        Nx.tensor([
+          [1.0, 2.0, 3.0],
+          [6, 10, 1]
+        ])
+
+      assert_equal(Nx.dot(a, [2], [0], Nx.LinAlg.triangular_solve(a, b), [1], [0]), b)
+    end
+
+    test "works with B that has more columns than rows" do
+      a =
+        Nx.tensor(
+          [
+            [1, 0],
+            [1, 1]
+          ],
+          type: :f64
+        )
+
+      b =
+        Nx.tensor(
+          [
+            [1, 1, 1],
+            [2, 2, 2]
+          ],
+          type: :f64
+        )
+
+      x = Nx.LinAlg.triangular_solve(a, b)
+
+      assert_equal(
+        x,
+        Nx.tensor(
+          [
+            [1, 1, 1],
+            [1, 1, 1]
+          ],
+          type: :f64
+        )
+      )
+    end
+
+    test "property" do
+      a = Nx.tensor([[1, 0, 0], [1, 1, 0], [0, 1, 1]])
+      b = Nx.tensor([[1.0, 2.0, 3.0], [2.0, 2.0, 4.0], [2.0, 0.0, 1.0]])
+      assert_equal(Nx.dot(a, Nx.LinAlg.triangular_solve(a, b)), b)
+
+      upper = Nx.transpose(a)
+      assert_equal(Nx.dot(upper, Nx.LinAlg.triangular_solve(upper, b, lower: false)), b)
+
+      assert_equal(
+        Nx.dot(
+          Nx.LinAlg.triangular_solve(upper, b, left_side: false, lower: false),
+          upper
+        ),
+        b
+      )
+
+      assert_equal(
+        Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose),
+        Nx.LinAlg.triangular_solve(upper, b, lower: false)
+      )
+
+      assert_equal(
+        Nx.dot(
+          Nx.transpose(a),
+          Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose)
+        ),
+        b
+      )
+    end
+  end
 end