Implement KroneckerProduct dispatch in mlx backend

Author: jessegrabowski

pytensor/link/mlx/dispatch/nlinalg.py

@@ -1,9 +1,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
-from pytensor.tensor.nlinalg import (
-    SVD,
-)
+from pytensor.tensor.nlinalg import SVD, KroneckerProduct
 
 
 @mlx_funcify.register(SVD)
@@ -35,3 +33,15 @@ def svd_full(x):
         return svd_full
     else:
         return svd_S_only
+
+
+@mlx_funcify.register(KroneckerProduct)
+def mlx_funcify_KroneckerProduct(op, node, **kwargs):
+    otype = node.outputs[0].dtype
+    mx_otype = getattr(mx, otype)
+    stream = mx.cpu if otype == "float64" else mx.gpu
+
+    def kron(a, b):
+        return mx.kron(a, b, stream=stream).astype(mx_otype, stream=stream)
+
+    return kron

tests/link/mlx/test_nlinalg.py

@@ -25,3 +25,19 @@ def test_mlx_svd(compute_uv):
         [A_val],
         mlx_mode=mlx_linalg_mode,
     )
+
+
+def test_mlx_kron():
+    rng = np.random.default_rng()
+
+    A = pt.matrix(name="A")
+    B = pt.matrix(name="B")
+    A_val, B_val = rng.normal(size=(2, 3, 3)).astype(config.floatX)
+    out = pt.linalg.kron(A, B)
+
+    compare_mlx_and_py(
+        [A, B],
+        [out],
+        [A_val, B_val],
+        mlx_mode=mlx_linalg_mode,
+    )