Implement SVD dispatch in mlx backend

Author: jessegrabowski

pytensor/link/mlx/dispatch/__init__.py

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

pytensor/link/mlx/dispatch/nlinalg.py

@@ -0,0 +1,37 @@
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.tensor.nlinalg import (
+    SVD,
+)
+
+
+@mlx_funcify.register(SVD)
+def mlx_funcify_SVD(op, node, **kwargs):
+    full_matrices = op.full_matrices
+    compute_uv = op.compute_uv
+    otype = (
+        getattr(mx, node.outputs[0].dtype)
+        if not compute_uv
+        else [getattr(mx, output.dtype) for output in node.outputs]
+    )
+
+    if not full_matrices:
+        raise TypeError("full_matrices=False is not supported in the mlx backend.")
+
+    def svd_S_only(x):
+        return mx.linalg.svd(x, compute_uv=False, stream=mx.cpu).astype(
+            otype, stream=mx.cpu
+        )
+
+    def svd_full(x):
+        outputs = mx.linalg.svd(x, compute_uv=True, stream=mx.cpu)
+        return (
+            output.astype(typ, stream=mx.cpu)
+            for output, typ in zip(outputs, otype, strict=True)
+        )
+
+    if compute_uv:
+        return svd_full
+    else:
+        return svd_S_only

tests/link/mlx/test_nlinalg.py

@@ -0,0 +1,27 @@
+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_linalg_mode = mlx_mode.including("blockwise")
+
+
+@pytest.mark.parametrize("compute_uv", [True, False])
+def test_mlx_svd(compute_uv):
+    rng = np.random.default_rng()
+
+    A = pt.matrix(name="X")
+    A_val = rng.normal(size=(3, 3)).astype(config.floatX)
+    A_val = A_val @ A_val.T
+
+    out = pt.linalg.svd(A, compute_uv=compute_uv)
+
+    compare_mlx_and_py(
+        [A],
+        out,
+        [A_val],
+        mlx_mode=mlx_linalg_mode,
+    )