Implement MatrixInv and MatrixPinv dispatch in mlx backend

Author: jessegrabowski

pytensor/link/mlx/dispatch/nlinalg.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
-from pytensor.tensor.nlinalg import SVD, KroneckerProduct
+from pytensor.tensor.nlinalg import SVD, KroneckerProduct, MatrixInverse, MatrixPinv
 
 
 @mlx_funcify.register(SVD)
@@ -45,3 +45,23 @@ def kron(a, b):
         return mx.kron(a, b, stream=stream).astype(mx_otype, stream=stream)
 
     return kron
+
+
+@mlx_funcify.register(MatrixInverse)
+def mlx_funcify_MatrixInverse(op, node, **kwargs):
+    otype = getattr(mx, node.outputs[0].dtype)
+
+    def inv(x):
+        return mx.linalg.inv(x, stream=mx.cpu).astype(otype, stream=mx.cpu)
+
+    return inv
+
+
+@mlx_funcify.register(MatrixPinv)
+def mlx_funcify_MatrixPinv(op, node, **kwargs):
+    otype = getattr(mx, node.outputs[0].dtype)
+
+    def pinv(x):
+        return mx.linalg.pinv(x, stream=mx.cpu).astype(otype, stream=mx.cpu)
+
+    return pinv

tests/link/mlx/test_nlinalg.py

@@ -41,3 +41,21 @@ def test_mlx_kron():
         [A_val, B_val],
         mlx_mode=mlx_linalg_mode,
     )
+
+
+@pytest.mark.parametrize("op", [pt.linalg.inv, pt.linalg.pinv], ids=["inv", "pinv"])
+def test_mlx_inv(op):
+    rng = np.random.default_rng()
+
+    A = pt.matrix(name="A")
+    A_val = rng.normal(size=(3, 3)).astype(config.floatX)
+    A_val = A_val @ A_val.T
+
+    out = op(A)
+
+    compare_mlx_and_py(
+        [A],
+        [out],
+        [A_val],
+        mlx_mode=mlx_linalg_mode,
+    )