Implement Solve dispatch in mlx backend

Author: jessegrabowski

pytensor/link/mlx/dispatch/__init__.py

@@ -10,4 +10,5 @@
 import pytensor.link.mlx.dispatch.signal
 import pytensor.link.mlx.dispatch.signal.conv
 import pytensor.link.mlx.dispatch.blockwise
+import pytensor.link.mlx.dispatch.slinalg
 # isort: on

pytensor/link/mlx/dispatch/slinalg.py

@@ -0,0 +1,25 @@
+import warnings
+
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.tensor.slinalg import Solve
+
+
+@mlx_funcify.register(Solve)
+def mlx_funcify_Solve(op, node, **kwargs):
+    assume_a = op.assume_a
+    out_dtype = getattr(mx, node.outputs[0].dtype)
+
+    if assume_a != "gen":
+        warnings.warn(
+            f"MLX solve does not support assume_a={op.assume_a}. Defaulting to assume_a='gen'.\n"
+            f"If appropriate, you may want to set assume_a to one of 'sym', 'pos', 'her' or 'tridiagonal' to improve performance.",
+            UserWarning,
+        )
+
+    def solve(a, b):
+        # MLX only supports solve on CPU
+        return mx.linalg.solve(a, b, stream=mx.cpu).astype(out_dtype, stream=mx.cpu)
+
+    return solve

tests/link/mlx/__init__.py

@@ -1,10 +1,11 @@
+import mlx as mx
 import numpy as np
 import pytest
 
 import pytensor
 import pytensor.tensor as pt
 from pytensor.tensor.math import Argmax, Max
-from tests.link.mlx.test_basic import compare_mlx_and_py, mx
+from tests.link.mlx.test_basic import compare_mlx_and_py
 
 
 def test_dot():

tests/link/mlx/test_math.py

@@ -0,0 +1,57 @@
+import contextlib
+
+import numpy as np
+import pytest
+
+import pytensor.tensor as pt
+from pytensor import config
+from tests.link.mlx.test_basic import compare_mlx_and_py, mlx_mode
+
+
+# mlx complains about useless vmap (when there are no batch dims), so we need to include
+# local_remove_useless_blockwise rewrite for these tests
+mlx_linalg_mode = mlx_mode.including("blockwise")
+
+
+@pytest.mark.parametrize("lower", [True, False])
+def test_mlx_cholesky(lower):
+    rng = np.random.default_rng()
+
+    A = pt.tensor("A", shape=(5, 5))
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    A_val = A_val @ A_val.T
+
+    out = pt.linalg.cholesky(A, lower=lower)
+
+    compare_mlx_and_py(
+        [A],
+        [out],
+        [A_val],
+        mlx_mode=mlx_linalg_mode,
+    )
+
+
+@pytest.mark.parametrize("assume_a", ["gen", "pos"])
+def test_mlx_solve(assume_a):
+    rng = np.random.default_rng()
+
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("B", shape=(5, 5))
+
+    out = pt.linalg.solve(A, b, b_ndim=2, assume_a=assume_a)
+
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    A_val = A_val @ A_val.T
+
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
+
+    context = (
+        contextlib.suppress()
+        if assume_a == "gen"
+        else pytest.warns(
+            UserWarning, match=f"MLX solve does not support assume_a={assume_a}"
+        )
+    )
+
+    with context:
+        compare_mlx_and_py([A, b], [out], [A_val, b_val], mlx_mode=mlx_linalg_mode)