Bypass mx.vmap for natively batched MLX linalg ops

[guillaume-osmo]

Closes #2385.

Stacked on #2396 (the batched triangular solve needs unit_diagonal to be honoured). Review that one first; this diff will shrink once it merges.

Motivation

Blockwise is lowered with mx.vmap, which is not a safe wrapper for every MLX linalg primitive. Two separate problems:

1. Missing batching rule (the reported crash). mx.linalg.solve and mx.linalg.lu_factor build an LUF primitive that has no vmap rule, so batched solve, lu_factor and det all raise [Primitive::vmap] Not implemented for LUF. (det is not mentioned in the issue but goes through lu_factor.)

2. mx.vmap's solve_triangular rule drops the upper flag. This one does not raise — it silently solves the full system:

import mlx.core as mx, numpy as np, scipy.linalg as sla
rng = np.random.default_rng(0)
A = rng.normal(size=(3, 3)).astype("float32")
b = rng.normal(size=(3,)).astype("float32")
Ab, bb = np.broadcast_to(A, (4, 3, 3)).copy(), np.broadcast_to(b, (4, 3)).copy()

print(sla.solve_triangular(A, b, lower=True))          # [-10.06  -0.807 -19.79]
print(np.asarray(mx.vmap(                              # [ -0.793 -0.256  -1.873]
    lambda a, y: mx.linalg.solve_triangular(a, y, upper=False, stream=mx.cpu)
)(mx.array(Ab), mx.array(bb)))[0])
print(np.linalg.solve(A, b))                           # [ -0.793 -0.256  -1.873]  <- matches vmap

The vmapped answer is the full solve, exactly. It is only correct when the input happens to be exactly triangular — which cho_solve and lu_solve never are, since lu_factor packs L and U into a single array. Wrong numbers, no warning; I consider this the more serious half.

cholesky, inv and PivotToPermutations vmap correctly, so this is specific to solve_triangular.

Implementation

MLX's own CPU-stream implementations already accept batched input, so the vmap wrapper is not needed for these in the first place. Add a mlx_funcify_batched singledispatch registry: a core Op may supply a natively batched implementation, and funcify_Blockwise then skips vmap entirely.

Registered for Solve, SolveTriangular, CholeskySolve, LUFactor, Det and SLogDet. _lu_det_parts now reduces over the trailing axes only, so it accepts a stack of matrices.

Two MLX details the registry has to absorb:

• mx.linalg.solve does not broadcast batch dims (unlike solve_triangular); it requires them to match. Inputs are aligned to a common batch shape first.
• MLX only special-cases a 1-d b. With a batch it rejects A (*batch, m, m) against b (*batch, m), so a vector right-hand side gets an explicit trailing column axis, removed again on the way out.

PivotToPermutations is deliberately left on the vmap path, since its Python loop is bounded by the core length there and is already correct; a test pins that so this change cannot regress it.

Tests

Verified against the reference for batch shapes (), (4,) and (2, 3) across solve (b_ndim 1 and 2), solve_triangular (both triangles, with and without unit_diagonal, on deliberately dense input), cho_solve, lu_solve, lu_factor, det, cholesky and inv, plus batch-dim broadcasting and gradients.

New tests, all of which fail without this change:

• test_mlx_solve_batched, test_mlx_solve_batch_broadcast
• test_mlx_SolveTriangular_batched_ignores_other_triangle — pins the silent corruption above
• test_mlx_lu_solve_batched, test_mlx_lu_factor_batched, test_mlx_det_batched
• test_mlx_pivot_to_permutations_batched — passes before and after, pinning the vmap path that is intentionally kept

Full tests/link/mlx/ suite: no regressions.

Note

Every mx.linalg.* call in this backend runs on stream=mx.cpu, so none of this is GPU work — relevant to #2383.