Implement SolveTriangular dispatch in mlx backend

Author: jessegrabowski

pytensor/link/mlx/dispatch/slinalg.py

@@ -3,7 +3,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
-from pytensor.tensor.slinalg import Solve
+from pytensor.tensor.slinalg import Solve, SolveTriangular
 
 
 @mlx_funcify.register(Solve)
@@ -23,3 +23,19 @@ def solve(a, b):
         return mx.linalg.solve(a, b, stream=mx.cpu).astype(out_dtype, stream=mx.cpu)
 
     return solve
+
+
+@mlx_funcify.register(SolveTriangular)
+def mlx_funcify_SolveTriangular(op, node, **kwargs):
+    lower = op.lower
+    out_dtype = getattr(mx, node.outputs[0].dtype)
+
+    def solve_triangular(A, b):
+        return mx.linalg.solve_triangular(
+            A,
+            b,
+            upper=not lower,
+            stream=mx.cpu,  # MLX only supports solve_triangular on CPU
+        ).astype(out_dtype, stream=mx.cpu)
+
+    return solve_triangular

tests/link/mlx/test_slinalg.py

@@ -55,3 +55,23 @@ def test_mlx_solve(assume_a):
 
     with context:
         compare_mlx_and_py([A, b], [out], [A_val, b_val], mlx_mode=mlx_linalg_mode)
+
+
+@pytest.mark.parametrize("lower, trans", [(False, False), (True, True)])
+def test_mlx_SolveTriangular(lower, trans):
+    rng = np.random.default_rng()
+
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("B", shape=(5, 5))
+
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
+
+    out = pt.linalg.solve_triangular(
+        A,
+        b,
+        trans=0,
+        lower=lower,
+        unit_diagonal=False,
+    )
+    compare_mlx_and_py([A, b], [out], [A_val, b_val], mlx_mode=mlx_linalg_mode)