Call MLX linalg ops directly on batched inputs instead of mx.vmap

pytensor/link/mlx/dispatch/blockwise.py

@@ -1,4 +1,5 @@
 import mlx.core as mx
+import numpy as np
 
 from pytensor.link.mlx.dispatch import mlx_funcify
 from pytensor.tensor.blockwise import Blockwise, _check_runtime_broadcast_core
@@ -19,6 +20,29 @@ def funcify_Blockwise(op: Blockwise, node, **kwargs):
     # Hoisted out of the per-call path, unlike Blockwise._check_runtime_broadcast.
     batch_bcast = [inp.type.broadcastable[:batch_ndim] for inp in node.inputs]
 
+    # A core function that sets `natively_batched` runs directly on inputs
+    # broadcast to the common batch shape. MLX's linalg functions take leading
+    # batch dims themselves, and mx.vmap has no rule for LUF or QRF and drops the
+    # triangular flags of solve_triangular.
+    if getattr(core_f, "natively_batched", False):
+
+        def blockwise_native(*args):
+            _check_runtime_broadcast_core(args, batch_bcast, batch_ndim)
+
+            batch_shapes = [
+                arg.shape[: arg.ndim - n_core] for arg, n_core in zip(args, core_ndims)
+            ]
+            batch_shape = np.broadcast_shapes(*batch_shapes)
+            args = [
+                mx.broadcast_to(arg, (*batch_shape, *arg.shape[arg.ndim - n_core :]))
+                for arg, n_core in zip(args, core_ndims)
+            ]
+
+            out = core_f(*args)
+            return tuple(out) if multi_output else out
+
+        return blockwise_native
+
     # Decide batching purely from static shapes so a graph batches identically
     # here and in every other backend: a batch axis broadcasts (is never mapped)
     # only when its static size is exactly 1, or the input lacks it entirely.

pytensor/link/mlx/dispatch/linalg/decomposition.py

@@ -29,10 +29,9 @@ def svd_full(x):
         )
         return outputs
 
-    if compute_uv:
-        return svd_full
-    else:
-        return svd_S_only
+    svd = svd_full if compute_uv else svd_S_only
+    svd.natively_batched = True
+    return svd
 
 
 @mlx_funcify.register(Cholesky)
@@ -45,6 +44,7 @@ def cholesky(a):
             a.astype(dtype=a_dtype, stream=mx.cpu), upper=not lower, stream=mx.cpu
         )
 
+    cholesky.natively_batched = True
     return cholesky
 
 
@@ -70,6 +70,7 @@ def lu(a):
             U,
         )
 
+    lu.natively_batched = True
     return lu
 
 
@@ -80,6 +81,7 @@ def mlx_funcify_Eig(op, node, **kwargs):
     def eig(x):
         return mx.linalg.eig(x.astype(dtype=X_dtype, stream=mx.cpu), stream=mx.cpu)
 
+    eig.natively_batched = True
     return eig
 
 
@@ -99,6 +101,7 @@ def eigh(a):
             a.astype(dtype=X_dtype, stream=mx.cpu), UPLO=UPLO, stream=mx.cpu
         )
 
+    eigh.natively_batched = True
     return eigh
 
 
@@ -118,6 +121,7 @@ def eigvalsh(a):
             a.astype(dtype=X_dtype, stream=mx.cpu), UPLO=UPLO, stream=mx.cpu
         )
 
+    eigvalsh.natively_batched = True
     return eigvalsh
 
 
@@ -131,6 +135,7 @@ def lu_factor(a):
         )
         return lu, pivots.astype(mx.int32, stream=mx.cpu)
 
+    lu_factor.natively_batched = True
     return lu_factor
 
 
@@ -174,4 +179,5 @@ def qr(a):
             return R
         return Q, R
 
+    qr.natively_batched = True
     return qr

pytensor/link/mlx/dispatch/linalg/inverse.py

@@ -11,6 +11,7 @@ def mlx_funcify_MatrixInverse(op, node, **kwargs):
     def inv(x):
         return mx.linalg.inv(x.astype(dtype=X_dtype, stream=mx.cpu), stream=mx.cpu)
 
+    inv.natively_batched = True
     return inv
 
 
@@ -21,4 +22,5 @@ def mlx_funcify_MatrixPinv(op, node, **kwargs):
     def pinv(x):
         return mx.linalg.pinv(x.astype(dtype=x_dtype, stream=mx.cpu), stream=mx.cpu)
 
+    pinv.natively_batched = True
     return pinv

pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -8,6 +8,15 @@
 from pytensor.tensor.linalg.solvers.triangular import SolveTriangular
 
 
+def _as_column(b, b_ndim):
+    # MLX treats a 2-d rhs as a matrix, so a batched vector rhs needs an explicit column.
+    return mx.expand_dims(b, -1, stream=mx.cpu) if b_ndim == 1 else b
+
+
+def _from_column(out, b_ndim):
+    return mx.squeeze(out, -1, stream=mx.cpu) if b_ndim == 1 else out
+
+
 @mlx_funcify.register(Solve)
 def mlx_funcify_Solve(op, node, **kwargs):
     assume_a = op.assume_a
@@ -20,13 +29,17 @@ def mlx_funcify_Solve(op, node, **kwargs):
             UserWarning,
         )
 
+    b_ndim = op.b_ndim
+
     def solve(a, b):
-        return mx.linalg.solve(
+        out = mx.linalg.solve(
             a.astype(stream=mx.cpu, dtype=a_dtype),
-            b.astype(stream=mx.cpu, dtype=b_dtype),
+            _as_column(b.astype(stream=mx.cpu, dtype=b_dtype), b_ndim),
             stream=mx.cpu,
         )
+        return _from_column(out, b_ndim)
 
+    solve.natively_batched = True
     return solve
 
 
@@ -36,6 +49,7 @@ def mlx_funcify_SolveTriangular(op, node, **kwargs):
     unit_diagonal = op.unit_diagonal
     A_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_ndim = op.b_ndim
 
     def solve_triangular(A, b):
         A = A.astype(stream=mx.cpu, dtype=A_dtype)
@@ -47,13 +61,15 @@ def solve_triangular(A, b):
             diagonal_mask = mx.eye(A.shape[-1], dtype=mx.bool_, stream=mx.cpu)
             A = mx.where(diagonal_mask, mx.array(1, dtype=A_dtype), A, stream=mx.cpu)
 
-        return mx.linalg.solve_triangular(
+        out = mx.linalg.solve_triangular(
             A,
-            b.astype(stream=mx.cpu, dtype=b_dtype),
+            _as_column(b.astype(stream=mx.cpu, dtype=b_dtype), b_ndim),
             upper=not lower,
             stream=mx.cpu,
         )
+        return _from_column(out, b_ndim)
 
+    solve_triangular.natively_batched = True
     return solve_triangular
 
 
@@ -62,15 +78,18 @@ def mlx_funcify_CholeskySolve(op, node, **kwargs):
     lower = op.lower
     c_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_ndim = op.b_ndim
 
     # MLX has no cho_solve, so with A = L L.T we solve L y = b then L.T x = y.
     def cho_solve(c, b):
         c = c.astype(stream=mx.cpu, dtype=c_dtype)
-        b = b.astype(stream=mx.cpu, dtype=b_dtype)
+        b = _as_column(b.astype(stream=mx.cpu, dtype=b_dtype), b_ndim)
         c_T = mx.swapaxes(c, -1, -2, stream=mx.cpu)
         L, L_T = (c, c_T) if lower else (c_T, c)
 
         y = mx.linalg.solve_triangular(L, b, upper=False, stream=mx.cpu)
-        return mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
+        out = mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
+        return _from_column(out, b_ndim)
 
+    cho_solve.natively_batched = True
     return cho_solve

pytensor/link/mlx/dispatch/linalg/summary.py

@@ -7,11 +7,11 @@
 def _lu_det_parts(x):
     """Compute sign and logdet via LU factorization. Call within a CPU stream context."""
     lu, pivots = mx.linalg.lu_factor(x)
-    diag_u = mx.diagonal(lu)
-    n_swaps = mx.sum(pivots != mx.arange(pivots.shape[0], dtype=pivots.dtype))
+    diag_u = mx.diagonal(lu, axis1=-2, axis2=-1)
+    n_swaps = mx.sum(pivots != mx.arange(pivots.shape[-1], dtype=pivots.dtype), axis=-1)
     pivot_sign = 1 - 2 * (n_swaps % 2)
-    sign = pivot_sign * mx.prod(mx.sign(diag_u))
-    logabsdet = mx.sum(mx.log(mx.abs(diag_u)))
+    sign = pivot_sign * mx.prod(mx.sign(diag_u), axis=-1)
+    logabsdet = mx.sum(mx.log(mx.abs(diag_u)), axis=-1)
     return sign, logabsdet
 
 
@@ -24,6 +24,7 @@ def det(x):
             sign, logabsdet = _lu_det_parts(x.astype(dtype=X_dtype))
             return sign * mx.exp(logabsdet)
 
+    det.natively_batched = True
     return det
 
 
@@ -35,4 +36,5 @@ def slogdet(x):
         with mx.stream(mx.cpu):
             return _lu_det_parts(x.astype(dtype=X_dtype))
 
+    slogdet.natively_batched = True
     return slogdet

tests/link/mlx/linalg/test_decomposition.py

@@ -83,13 +83,14 @@ def test_mlx_cholesky(lower):
     )
 
 
-def test_mlx_LU():
+@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"])
+def test_mlx_LU(batch_shape):
     rng = np.random.default_rng(15)
 
-    A = pt.tensor("A", shape=(5, 5))
+    A = pt.tensor("A", shape=(*batch_shape, 5, 5))
     out = lu.lu(A, permute_l=False, p_indices=True)
 
-    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    A_val = rng.normal(size=(*batch_shape, 5, 5)).astype(config.floatX)
 
     compare_mlx_and_py(
         [A],
@@ -120,11 +121,12 @@ def test_mlx_eigvalsh(lower):
     compare_mlx_and_py([A], [out_no_b], [A_val])
 
 
-def test_mlx_lu_factor():
+@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"])
+def test_mlx_lu_factor(batch_shape):
     rng = np.random.default_rng(15)
 
-    A = pt.matrix(name="A")
-    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    A = pt.tensor("A", shape=(*batch_shape, 5, 5))
+    A_val = rng.normal(size=(*batch_shape, 5, 5)).astype(config.floatX)
 
     out = pt.linalg.lu_factor(A)
 
@@ -166,12 +168,13 @@ def test_mlx_pivot_to_permutations():
     compare_mlx_and_py([A], [out], [A_val])
 
 
+@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"])
 @pytest.mark.parametrize("mode", ["economic", "r"])
-def test_mlx_qr(mode):
+def test_mlx_qr(mode, batch_shape):
     rng = np.random.default_rng(15)
 
-    A = pt.matrix(name="A")
-    A_val = rng.normal(size=(5, 3)).astype(config.floatX)
+    A = pt.tensor("A", shape=(*batch_shape, 5, 3))
+    A_val = rng.normal(size=(*batch_shape, 5, 3)).astype(config.floatX)
 
     out = pt.linalg.qr(A, mode=mode)
 

tests/link/mlx/linalg/test_inverse.py

@@ -8,14 +8,15 @@
 from tests.link.mlx.test_basic import compare_mlx_and_py, mlx_mode
 
 
+@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"])
 @pytest.mark.parametrize("op", [pt.linalg.inv, pt.linalg.pinv], ids=["inv", "pinv"])
-def test_mlx_inv(op):
+def test_mlx_inv(op, batch_shape):
     rng = np.random.default_rng(15)
     n = 3
 
-    A = pt.matrix(name="A")
-    A_val = rng.normal(size=(n, n))
-    A_val = (A_val @ A_val.T).astype(config.floatX)
+    A = pt.tensor("A", shape=(*batch_shape, n, n))
+    A_val = rng.normal(size=(*batch_shape, n, n))
+    A_val = (A_val @ np.swapaxes(A_val, -1, -2)).astype(config.floatX)
 
     out = op(A)
 

tests/link/mlx/linalg/test_solvers.py

@@ -9,20 +9,21 @@
 from tests.link.mlx.test_basic import compare_mlx_and_py, mlx_mode
 
 
+@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"])
 @pytest.mark.parametrize("assume_a", ["gen", "pos"])
-def test_mlx_solve(assume_a):
+def test_mlx_solve(assume_a, batch_shape):
     rng = np.random.default_rng(15)
     n = 3
 
-    A = pt.tensor("A", shape=(n, n))
-    b = pt.tensor("B", shape=(n, n))
+    A = pt.tensor("A", shape=(*batch_shape, n, n))
+    b = pt.tensor("B", shape=(*batch_shape, n, n))
 
     out = pt.linalg.solve(A, b, b_ndim=2, assume_a=assume_a)
 
-    A_val = rng.normal(size=(n, n)).astype(config.floatX)
-    A_val = A_val @ A_val.T
+    A_val = rng.normal(size=(*batch_shape, n, n)).astype(config.floatX)
+    A_val = A_val @ np.swapaxes(A_val, -1, -2)
 
-    b_val = rng.normal(size=(n, n)).astype(config.floatX)
+    b_val = rng.normal(size=(*batch_shape, n, n)).astype(config.floatX)
 
     context = (
         contextlib.suppress()
@@ -44,20 +45,24 @@ def test_mlx_solve(assume_a):
         )
 
 
+@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"])
 @pytest.mark.parametrize(
     "unit_diagonal", [False, True], ids=["full_diagonal", "unit_diagonal"]
 )
 @pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"])
-def test_mlx_SolveTriangular(lower, unit_diagonal):
+def test_mlx_SolveTriangular(lower, unit_diagonal, batch_shape):
     rng = np.random.default_rng(15)
 
-    A = pt.tensor("A", shape=(5, 5))
-    b = pt.tensor("B", shape=(5, 5))
+    A = pt.tensor("A", shape=(*batch_shape, 5, 5))
+    b = pt.tensor("B", shape=(*batch_shape, 5, 5))
 
-    # A diagonal far from one, so ignoring `unit_diagonal` gives a different answer
-    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
-    A_val[np.diag_indices(5)] = rng.uniform(3, 4, size=5).astype(config.floatX)
-    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    # A dense matrix with a diagonal far from one, so both ignoring the
+    # triangle and ignoring `unit_diagonal` give a different answer
+    A_val = rng.normal(size=(*batch_shape, 5, 5)).astype(config.floatX)
+    A_val[..., np.arange(5), np.arange(5)] = rng.uniform(
+        3, 4, size=(*batch_shape, 5)
+    ).astype(config.floatX)
+    b_val = rng.normal(size=(*batch_shape, 5, 5)).astype(config.floatX)
 
     out = pt.linalg.solve_triangular(
         A,
@@ -109,6 +114,34 @@ def test_mlx_CholeskySolve(batch_shape, lower, b_ndim):
     )
 
 
+@pytest.mark.parametrize(
+    "a_batch, b_batch",
+    [((4,), ()), ((2, 1), (1, 3))],
+    ids=["unbatched_rhs", "cross_broadcast"],
+)
+def test_mlx_solve_batch_broadcasting(a_batch, b_batch):
+    rng = np.random.default_rng(15)
+    n = 3
+
+    A = pt.tensor("A", shape=(*a_batch, n, n))
+    b = pt.tensor("b", shape=(*b_batch, n))
+    out = pt.linalg.solve(A, b, b_ndim=1)
+
+    A_val = rng.normal(size=(*a_batch, n, n)).astype(config.floatX)
+    A_val = A_val @ np.swapaxes(A_val, -1, -2) + n * np.eye(n, dtype=config.floatX)
+    b_val = rng.normal(size=(*b_batch, n)).astype(config.floatX)
+
+    compare_mlx_and_py(
+        [A, b],
+        [out],
+        [A_val, b_val],
+        mlx_mode=mlx_mode,
+        assert_fn=partial(
+            np.testing.assert_allclose, atol=1e-6, rtol=1e-6, strict=True
+        ),
+    )
+
+
 def test_mlx_CholeskySolve_mixed_dtypes():
     rng = np.random.default_rng(15)
     n = 5