Bypass mx.vmap for natively batched MLX linalg ops

pytensor/link/mlx/dispatch/blockwise.py

@@ -1,9 +1,46 @@
+from functools import singledispatch
+
 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
 
 
+@singledispatch
+def mlx_funcify_batched(core_op, node, **kwargs):
+    """Return a natively batched implementation of ``core_op``, or ``None``.
+
+    ``Blockwise`` is lowered with ``mx.vmap``, which has no batching rule for
+    every MLX primitive: ``mx.linalg.solve`` and ``mx.linalg.lu_factor`` both
+    build an ``LUF`` primitive and raise "[Primitive::vmap] Not implemented for
+    LUF" (#2385). Their CPU-stream implementations already accept batched
+    input, so the ``vmap`` wrapper is not needed for them in the first place.
+
+    An ``Op`` registered here receives all of its inputs already aligned to a
+    common batch shape -- ``mx.linalg.solve`` requires identical batch dims
+    across inputs and, unlike ``solve_triangular``, does not broadcast them
+    itself.
+    """
+    return None
+
+
+def _align_batch_dims(args, core_ndims):
+    """Broadcast every input to the common batch shape implied by ``core_ndims``."""
+    batch_shapes = [
+        arg.shape[: arg.ndim - n_core] for arg, n_core in zip(args, core_ndims)
+    ]
+    batch_shape = tuple(np.broadcast_shapes(*batch_shapes))
+
+    aligned = []
+    for arg, n_core in zip(args, core_ndims):
+        target = (*batch_shape, *arg.shape[arg.ndim - n_core :])
+        aligned.append(
+            arg if arg.shape == target else mx.broadcast_to(arg, target, stream=mx.cpu)
+        )
+    return aligned
+
+
 @mlx_funcify.register(Blockwise)
 def funcify_Blockwise(op: Blockwise, node, **kwargs):
     core_node = op._create_dummy_core_node(node.inputs)
@@ -19,6 +56,19 @@ 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]
 
+    # Prefer a natively batched implementation over `mx.vmap` when the core `Op`
+    # provides one, both to sidestep missing `vmap` rules (#2385) and to let MLX
+    # batch the primitive itself.
+    batched_f = mlx_funcify_batched(op.core_op, node=core_node, **kwargs)
+    if batched_f is not None:
+
+        def blockwise_batched(*args):
+            _check_runtime_broadcast_core(args, batch_bcast, batch_ndim)
+            out = batched_f(*_align_batch_dims(args, core_ndims))
+            return tuple(out) if multi_output else out
+
+        return blockwise_batched
+
     # 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

@@ -1,6 +1,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.link.mlx.dispatch.blockwise import mlx_funcify_batched
 from pytensor.tensor.linalg.decomposition.cholesky import Cholesky
 from pytensor.tensor.linalg.decomposition.eigen import Eig, Eigh, Eigvalsh
 from pytensor.tensor.linalg.decomposition.lu import LU, LUFactor, PivotToPermutations
@@ -175,3 +176,16 @@ def qr(a):
         return Q, R
 
     return qr
+
+
+@mlx_funcify_batched.register(LUFactor)
+def mlx_funcify_batched_LUFactor(op, node, **kwargs):
+    """`mx.linalg.lu_factor` is already batched; `mx.vmap` has no `LUF` rule (#2385)."""
+    A_dtype = getattr(mx, node.inputs[0].dtype)
+
+    def lu_factor(a):
+        with mx.stream(mx.cpu):
+            lu, pivots = mx.linalg.lu_factor(a.astype(dtype=A_dtype))
+            return lu, pivots.astype(mx.int32)
+
+    return lu_factor

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

@@ -3,6 +3,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.link.mlx.dispatch.blockwise import mlx_funcify_batched
 from pytensor.tensor.linalg.solvers.general import Solve
 from pytensor.tensor.linalg.solvers.psd import CholeskySolve
 from pytensor.tensor.linalg.solvers.triangular import SolveTriangular
@@ -30,19 +31,32 @@ def solve(a, b):
     return solve
 
 
+def _unit_diagonal(A):
+    """Replace the stored diagonal of ``A`` with ones.
+
+    Call within a CPU stream context. Broadcasts over any leading batch dims.
+    """
+    eye = mx.eye(A.shape[-1], dtype=A.dtype)
+    return A * (1 - eye) + eye
+
+
 @mlx_funcify.register(SolveTriangular)
 def mlx_funcify_SolveTriangular(op, node, **kwargs):
     lower = op.lower
+    unit_diagonal = op.unit_diagonal
     A_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
 
     def solve_triangular(A, b):
-        return mx.linalg.solve_triangular(
-            A.astype(stream=mx.cpu, dtype=A_dtype),
-            b.astype(stream=mx.cpu, dtype=b_dtype),
-            upper=not lower,
-            stream=mx.cpu,
-        )
+        with mx.stream(mx.cpu):
+            A = A.astype(dtype=A_dtype)
+            b = b.astype(dtype=b_dtype)
+            if unit_diagonal:
+                # `mx.linalg.solve_triangular` has no `unit_diagonal` argument,
+                # so the stored diagonal would be used instead of ones and the
+                # wrong answer returned silently (#2384).
+                A = _unit_diagonal(A)
+            return mx.linalg.solve_triangular(A, b, upper=not lower)
 
     return solve_triangular
 
@@ -64,3 +78,85 @@ def cho_solve(c, b):
         return mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
 
     return cho_solve
+
+
+def _as_column(b, is_vector):
+    """Give a vector right-hand side an explicit trailing column axis.
+
+    MLX only special-cases a 1-d `b`; with a batch it rejects `A (*batch, m, m)`
+    against `b (*batch, m)`, so the column has to be materialised.
+    """
+    return mx.expand_dims(b, -1) if is_vector else b
+
+
+@mlx_funcify_batched.register(Solve)
+def mlx_funcify_batched_Solve(op, node, **kwargs):
+    """`mx.linalg.solve` is already batched, so `Blockwise` must not `vmap` it.
+
+    `mx.vmap` has no rule for the `LUF` primitive that `solve` builds, so every
+    batched solve raised "[Primitive::vmap] Not implemented for LUF" (#2385).
+    """
+    a_dtype = getattr(mx, node.inputs[0].dtype)
+    b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_is_vector = node.inputs[1].type.ndim == 1
+
+    def solve(a, b):
+        with mx.stream(mx.cpu):
+            out = mx.linalg.solve(
+                a.astype(dtype=a_dtype),
+                _as_column(b.astype(dtype=b_dtype), b_is_vector),
+            )
+            return out.squeeze(-1) if b_is_vector else out
+
+    return solve
+
+
+@mlx_funcify_batched.register(SolveTriangular)
+def mlx_funcify_batched_SolveTriangular(op, node, **kwargs):
+    """Bypass `mx.vmap`, whose `solve_triangular` rule drops the `upper` flag.
+
+    A batched `solve_triangular` did not raise -- it silently solved the *full*
+    system instead of the triangular one, so any input that was not already
+    exactly triangular returned wrong numbers. That is the normal case for
+    `cho_solve` and `lu_solve`, which pack `L` and `U` into one array. The
+    eager `mx.linalg.solve_triangular` batches correctly.
+    """
+    lower = op.lower
+    unit_diagonal = op.unit_diagonal
+    A_dtype = getattr(mx, node.inputs[0].dtype)
+    b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_is_vector = node.inputs[1].type.ndim == 1
+
+    def solve_triangular(A, b):
+        with mx.stream(mx.cpu):
+            A = A.astype(dtype=A_dtype)
+            if unit_diagonal:
+                A = _unit_diagonal(A)
+            out = mx.linalg.solve_triangular(
+                A, _as_column(b.astype(dtype=b_dtype), b_is_vector), upper=not lower
+            )
+            return out.squeeze(-1) if b_is_vector else out
+
+    return solve_triangular
+
+
+@mlx_funcify_batched.register(CholeskySolve)
+def mlx_funcify_batched_CholeskySolve(op, node, **kwargs):
+    """As for `SolveTriangular`: `cho_solve` is two triangular solves."""
+    lower = op.lower
+    c_dtype = getattr(mx, node.inputs[0].dtype)
+    b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_is_vector = node.inputs[1].type.ndim == 1
+
+    def cho_solve(c, b):
+        with mx.stream(mx.cpu):
+            c = c.astype(dtype=c_dtype)
+            c_T = mx.swapaxes(c, -1, -2)
+            L, L_T = (c, c_T) if lower else (c_T, c)
+
+            b = _as_column(b.astype(dtype=b_dtype), b_is_vector)
+            y = mx.linalg.solve_triangular(L, b, upper=False)
+            out = mx.linalg.solve_triangular(L_T, y, upper=True)
+            return out.squeeze(-1) if b_is_vector else out
+
+    return cho_solve

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

@@ -1,17 +1,22 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.link.mlx.dispatch.blockwise import mlx_funcify_batched
 from pytensor.tensor.linalg.summary import Det, SLogDet
 
 
 def _lu_det_parts(x):
-    """Compute sign and logdet via LU factorization. Call within a CPU stream context."""
+    """Compute sign and logdet via LU factorization. Call within a CPU stream context.
+
+    Reduces over the trailing axes only, so a stack of matrices is handled
+    directly -- `mx.linalg.lu_factor` is natively batched (#2385).
+    """
     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
 
 
@@ -36,3 +41,28 @@ def slogdet(x):
             return _lu_det_parts(x.astype(dtype=X_dtype))
 
     return slogdet
+
+
+@mlx_funcify_batched.register(Det)
+def mlx_funcify_batched_Det(op, node, **kwargs):
+    """`Det` goes through `mx.linalg.lu_factor`, which `mx.vmap` cannot batch (#2385)."""
+    X_dtype = getattr(mx, node.inputs[0].dtype)
+
+    def det(x):
+        with mx.stream(mx.cpu):
+            sign, logabsdet = _lu_det_parts(x.astype(dtype=X_dtype))
+            return sign * mx.exp(logabsdet)
+
+    return det
+
+
+@mlx_funcify_batched.register(SLogDet)
+def mlx_funcify_batched_SLogDet(op, node, **kwargs):
+    """As for `Det`."""
+    X_dtype = getattr(mx, node.inputs[0].dtype)
+
+    def slogdet(x):
+        with mx.stream(mx.cpu):
+            return _lu_det_parts(x.astype(dtype=X_dtype))
+
+    return slogdet

tests/link/mlx/linalg/test_decomposition.py

@@ -155,3 +155,49 @@ def test_mlx_qr(mode):
     out = pt.linalg.qr(A, mode=mode)
 
     compare_mlx_and_py([A], out, [A_val])
+
+
+@pytest.mark.parametrize("batch_shape", [(4,), (2, 3)], ids=["batch", "batch_2d"])
+def test_mlx_lu_factor_batched(batch_shape):
+    """`mx.vmap` has no rule for the `LUF` primitive `lu_factor` builds (#2385)."""
+    rng = np.random.default_rng(15)
+    n = 5
+
+    A = pt.tensor("A", shape=(*batch_shape, n, n))
+    A_val = rng.normal(size=(*batch_shape, n, n)).astype(config.floatX)
+
+    compare_mlx_and_py([A], pt.linalg.lu_factor(A), [A_val], mlx_mode=mlx_mode)
+
+
+@pytest.mark.parametrize("inverse", [True, False], ids=["inverse", "forward"])
+@pytest.mark.parametrize("batch_shape", [(4,), (2, 3)], ids=["batch", "batch_2d"])
+def test_mlx_pivot_to_permutations_batched(batch_shape, inverse):
+    """Coverage for the batched pivot scan reached by a batched `lu_solve`.
+
+    The core dispatch loops in Python over ``pivots.shape[0]``, which is the
+    *core* length under `mx.vmap`, so this path is already correct -- pinning it
+    so the batched `lu_solve` fix cannot regress it. Built directly because
+    `pivot_to_permutation` itself only accepts a 1-d input.
+    """
+    from pytensor.tensor.blockwise import Blockwise
+    from pytensor.tensor.linalg.decomposition.lu import PivotToPermutations
+
+    rng = np.random.default_rng(15)
+    n = 5
+
+    pivots = pt.tensor("pivots", shape=(*batch_shape, n), dtype="int32")
+    # LAPACK-style pivots: entry i may only reference row i or later.
+    pivots_val = (
+        np.stack(
+            [
+                np.array([rng.integers(i, n) for i in range(n)])
+                for _ in range(int(np.prod(batch_shape)))
+            ]
+        )
+        .reshape(*batch_shape, n)
+        .astype("int32")
+    )
+
+    out = Blockwise(PivotToPermutations(inverse=inverse))(pivots)
+
+    compare_mlx_and_py([pivots], [out], [pivots_val], mlx_mode=mlx_mode)