From 868ed4e2bb8002387ccdebc38b26eac43aca9f57 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Thu, 17 Sep 2026 23:21:22 +0800 Subject: [PATCH 1/4] MLX: call natively batched core ops directly in Blockwise Co-authored-by: guillaume-osmo --- pytensor/link/mlx/dispatch/blockwise.py | 24 ++++++++++++++++++++++++ tests/link/mlx/test_blockwise.py | 14 ++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/pytensor/link/mlx/dispatch/blockwise.py b/pytensor/link/mlx/dispatch/blockwise.py index fe4ce5d1d0..828aacaac4 100644 --- a/pytensor/link/mlx/dispatch/blockwise.py +++ b/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. diff --git a/tests/link/mlx/test_blockwise.py b/tests/link/mlx/test_blockwise.py index 891ca741c6..8b201143b2 100644 --- a/tests/link/mlx/test_blockwise.py +++ b/tests/link/mlx/test_blockwise.py @@ -105,6 +105,20 @@ def test_blockwise_no_runtime_broadcast(): mlx_fn(*values) +def test_blockwise_native_no_runtime_broadcast(): + rng = np.random.default_rng(7) + a = tensor("a", shape=(None, 3, 3)) + b = tensor("b", shape=(5, 3, 2)) + out = pt.linalg.solve(a, b) + + assert isinstance(out.owner.op, Blockwise) + values = [rng.standard_normal((1, 3, 3)), rng.standard_normal((5, 3, 2))] + + mlx_fn = pytensor.function([a, b], out, mode=mlx_mode) + with pytest.raises(ValueError, match="Runtime broadcasting not allowed"): + mlx_fn(*values) + + @pytest.mark.parametrize("batch", [(), (5,)], ids=["no_batch", "single_batch"]) def test_blockwise_fallback_signature(batch): rng = np.random.default_rng(7) From 8e77ce122190f4291c03065c30abcb2f1872c978 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Thu, 17 Sep 2026 23:21:22 +0800 Subject: [PATCH 2/4] MLX: batch the solve dispatches natively Co-authored-by: guillaume-osmo --- pytensor/link/mlx/dispatch/linalg/solvers.py | 31 ++++++++-- tests/link/mlx/linalg/test_solvers.py | 59 +++++++++++++++----- 2 files changed, 71 insertions(+), 19 deletions(-) diff --git a/pytensor/link/mlx/dispatch/linalg/solvers.py b/pytensor/link/mlx/dispatch/linalg/solvers.py index 67c35b20fc..70620885c3 100644 --- a/pytensor/link/mlx/dispatch/linalg/solvers.py +++ b/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 diff --git a/tests/link/mlx/linalg/test_solvers.py b/tests/link/mlx/linalg/test_solvers.py index 1a4e88d9cc..a60ddb0dbb 100644 --- a/tests/link/mlx/linalg/test_solvers.py +++ b/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 From 4fdd356c467e7033f7fb9db0b629e36e66ce0865 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Thu, 17 Sep 2026 23:21:22 +0800 Subject: [PATCH 3/4] MLX: batch det and slogdet natively Co-authored-by: guillaume-osmo --- pytensor/link/mlx/dispatch/linalg/summary.py | 10 ++++++---- tests/link/mlx/linalg/test_summary.py | 14 ++++++++------ 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/pytensor/link/mlx/dispatch/linalg/summary.py b/pytensor/link/mlx/dispatch/linalg/summary.py index ed6d6507c0..2fd3b9424c 100644 --- a/pytensor/link/mlx/dispatch/linalg/summary.py +++ b/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 diff --git a/tests/link/mlx/linalg/test_summary.py b/tests/link/mlx/linalg/test_summary.py index de587bf59b..8d85656ab9 100644 --- a/tests/link/mlx/linalg/test_summary.py +++ b/tests/link/mlx/linalg/test_summary.py @@ -10,22 +10,24 @@ mx = pytest.importorskip("mlx.core") -def test_mlx_det(): +@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"]) +def test_mlx_det(batch_shape): rng = np.random.default_rng(15) - A = pt.matrix(name="A") - A_val = rng.normal(size=(3, 3)).astype(config.floatX) + A = pt.tensor("A", shape=(*batch_shape, 3, 3)) + A_val = rng.normal(size=(*batch_shape, 3, 3)).astype(config.floatX) out = pt.linalg.det(A) compare_mlx_and_py([A], [out], [A_val]) -def test_mlx_slogdet(): +@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"]) +def test_mlx_slogdet(batch_shape): rng = np.random.default_rng(15) - A = pt.matrix(name="A") - A_val = rng.normal(size=(3, 3)).astype(config.floatX) + A = pt.tensor("A", shape=(*batch_shape, 3, 3)) + A_val = rng.normal(size=(*batch_shape, 3, 3)).astype(config.floatX) sign, logabsdet = pt.linalg.slogdet(A) From 4362114517551088a0a7573dbad2ba7215d94212 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Thu, 17 Sep 2026 23:21:23 +0800 Subject: [PATCH 4/4] MLX: batch decompositions and inverses natively --- .../link/mlx/dispatch/linalg/decomposition.py | 14 +++++++++---- pytensor/link/mlx/dispatch/linalg/inverse.py | 2 ++ tests/link/mlx/linalg/test_decomposition.py | 21 +++++++++++-------- tests/link/mlx/linalg/test_inverse.py | 9 ++++---- 4 files changed, 29 insertions(+), 17 deletions(-) diff --git a/pytensor/link/mlx/dispatch/linalg/decomposition.py b/pytensor/link/mlx/dispatch/linalg/decomposition.py index fbd770b841..a797874ad3 100644 --- a/pytensor/link/mlx/dispatch/linalg/decomposition.py +++ b/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 diff --git a/pytensor/link/mlx/dispatch/linalg/inverse.py b/pytensor/link/mlx/dispatch/linalg/inverse.py index a98c593521..48ae1d7707 100644 --- a/pytensor/link/mlx/dispatch/linalg/inverse.py +++ b/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 diff --git a/tests/link/mlx/linalg/test_decomposition.py b/tests/link/mlx/linalg/test_decomposition.py index 81295b6e68..7e9d35d380 100644 --- a/tests/link/mlx/linalg/test_decomposition.py +++ b/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) diff --git a/tests/link/mlx/linalg/test_inverse.py b/tests/link/mlx/linalg/test_inverse.py index b0d6d97f2b..a09cddf174 100644 --- a/tests/link/mlx/linalg/test_inverse.py +++ b/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)