fix(exla): triangular_solve with batched matrix input

Author: TomasPegado

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/lib/exla/mlir/value.ex

@@ -598,18 +598,7 @@ defmodule EXLA.MLIR.Value do
         transform,
         typespec
       ) do
-
-    new_b_shape = {2, 3, 1}
-    new_b_typespec = %{typespec | shape: new_b_shape}
-    b = reshape(b, new_b_typespec)
-
-    expected_output_shape = {2, 3, 1}
-    new_typespec = %{typespec | shape: expected_output_shape}
-    result_types = typespecs_to_mlir_types([new_typespec])
-
-
-
-    # result_types = typespecs_to_mlir_types([typespec])
+    result_types = typespecs_to_mlir_types([typespec])
 
     complex? = Nx.Type.complex?(typespec.type)
 
@@ -627,13 +616,8 @@ defmodule EXLA.MLIR.Value do
       transpose_a: transpose_a
     ]
 
-
-
     op(func, "stablehlo.triangular_solve", [a, b], result_types, attributes: attributes)
     |> one!()
-    |> reshape(%{typespec | shape: {2, 3}})
-
-
   end
 
   def dynamic_update_slice(%Value{function: func} = operand, updates, starts, typespec) do

exla/test/exla/nx_linalg_doctest_test.exs

@@ -20,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 ++
@@ -430,7 +431,7 @@ defmodule EXLA.NxLinAlgDoctestTest do
           [6, 10, 1]
         ])
 
-      assert Nx.dot(a, [2], [0], Nx.LinAlg.triangular_solve(a, b), [1], [0]) == b
+      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
@@ -454,36 +455,46 @@ defmodule EXLA.NxLinAlgDoctestTest do
 
       x = Nx.LinAlg.triangular_solve(a, b)
 
-      assert x ==
-               Nx.tensor(
-                 [
-                   [1, 1, 1],
-                   [1, 1, 1]
-                 ],
-                 type: :f64
-               )
+      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 Nx.dot(a, Nx.LinAlg.triangular_solve(a, b)) == b
+      assert_equal(Nx.dot(a, Nx.LinAlg.triangular_solve(a, b)), b)
 
       upper = Nx.transpose(a)
-      assert Nx.dot(upper, Nx.LinAlg.triangular_solve(upper, b, lower: false)) == b
-
-      assert Nx.dot(
-               Nx.LinAlg.triangular_solve(upper, b, left_side: false, lower: false),
-               upper
-             ) == b
+      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 Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose) ==
-               Nx.LinAlg.triangular_solve(upper, b, lower: false)
+      assert_equal(
+        Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose),
+        Nx.LinAlg.triangular_solve(upper, b, lower: false)
+      )
 
-      assert Nx.dot(
-               Nx.transpose(a),
-               Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose)
-             ) == b
+      assert_equal(
+        Nx.dot(
+          Nx.transpose(a),
+          Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose)
+        ),
+        b
+      )
     end
   end
 end