Honour unit_diagonal in the MLX SolveTriangular dispatch

pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -30,19 +30,32 @@ def solve(a, b):
     return solve
 
 
+def _unit_diagonal(A):
+    """Replace the stored diagonal of ``A`` with ones.
+
+    Call within a CPU stream context. Broadcasts over any leading batch dims.
+    """
+    eye = mx.eye(A.shape[-1], dtype=A.dtype)
+    return A * (1 - eye) + eye
+
+
 @mlx_funcify.register(SolveTriangular)
 def mlx_funcify_SolveTriangular(op, node, **kwargs):
     lower = op.lower
+    unit_diagonal = op.unit_diagonal
     A_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
 
     def solve_triangular(A, b):
-        return mx.linalg.solve_triangular(
-            A.astype(stream=mx.cpu, dtype=A_dtype),
-            b.astype(stream=mx.cpu, dtype=b_dtype),
-            upper=not lower,
-            stream=mx.cpu,
-        )
+        with mx.stream(mx.cpu):
+            A = A.astype(dtype=A_dtype)
+            b = b.astype(dtype=b_dtype)
+            if unit_diagonal:
+                # `mx.linalg.solve_triangular` has no `unit_diagonal` argument,
+                # so the stored diagonal would be used instead of ones and the
+                # wrong answer returned silently (#2384).
+                A = _unit_diagonal(A)
+            return mx.linalg.solve_triangular(A, b, upper=not lower)
 
     return solve_triangular
 

tests/link/mlx/linalg/test_solvers.py

@@ -44,8 +44,14 @@ def test_mlx_solve(assume_a):
         )
 
 
+@pytest.mark.parametrize(
+    "unit_diagonal", [False, True], ids=["stored_diagonal", "unit_diagonal"]
+)
 @pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"])
-def test_mlx_SolveTriangular(lower):
+def test_mlx_SolveTriangular(lower, unit_diagonal):
+    # `unit_diagonal=True` was dropped on the way to `mx.linalg.solve_triangular`,
+    # which has no such argument, so the stored diagonal was used instead of ones
+    # and a wrong answer was returned without raising (#2384).
     rng = np.random.default_rng(15)
 
     A = pt.tensor("A", shape=(5, 5))
@@ -59,7 +65,7 @@ def test_mlx_SolveTriangular(lower):
         b,
         trans=0,
         lower=lower,
-        unit_diagonal=False,
+        unit_diagonal=unit_diagonal,
     )
     compare_mlx_and_py(
         [A, b],
@@ -128,3 +134,32 @@ def test_mlx_CholeskySolve_mixed_dtypes():
             np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True
         ),
     )
+
+
+def test_mlx_SolveTriangular_unit_diagonal_lu_solve():
+    # Batched `lu_factor` is blocked on #2385, so this covers the core case only.
+    batch_shape = ()
+    # `lu_solve` packs a unit-triangular `L` into the LU array, with `U`'s
+    # diagonal occupying those slots, so it relies on `unit_diagonal=True`
+    # actually being honoured (#2384).
+    rng = np.random.default_rng(15)
+    n = 4
+
+    A = pt.tensor("A", shape=(*batch_shape, n, n))
+    b = pt.tensor("b", shape=(*batch_shape, n))
+
+    A_val = rng.normal(size=(*batch_shape, n, n)).astype(config.floatX)
+    A_val = A_val @ np.swapaxes(A_val, -1, -2) + n * np.eye(n, dtype=config.floatX)
+    b_val = rng.normal(size=(*batch_shape, n)).astype(config.floatX)
+
+    out = pt.linalg.lu_solve(pt.linalg.lu_factor(A), b, b_ndim=1)
+
+    compare_mlx_and_py(
+        [A, b],
+        [out],
+        [A_val, b_val],
+        mlx_mode=mlx_mode,
+        assert_fn=partial(
+            np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True
+        ),
+    )