diff --git a/pytensor/link/mlx/dispatch/blockwise.py b/pytensor/link/mlx/dispatch/blockwise.py index fe4ce5d1d0..480807a290 100644 --- a/pytensor/link/mlx/dispatch/blockwise.py +++ b/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. diff --git a/pytensor/link/mlx/dispatch/linalg/decomposition.py b/pytensor/link/mlx/dispatch/linalg/decomposition.py index fbd770b841..2d23ce75ae 100644 --- a/pytensor/link/mlx/dispatch/linalg/decomposition.py +++ b/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 diff --git a/pytensor/link/mlx/dispatch/linalg/solvers.py b/pytensor/link/mlx/dispatch/linalg/solvers.py index 2f3f31416a..9b8a451bdb 100644 --- a/pytensor/link/mlx/dispatch/linalg/solvers.py +++ b/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 diff --git a/pytensor/link/mlx/dispatch/linalg/summary.py b/pytensor/link/mlx/dispatch/linalg/summary.py index ed6d6507c0..ab5b562ab3 100644 --- a/pytensor/link/mlx/dispatch/linalg/summary.py +++ b/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 diff --git a/tests/link/mlx/linalg/test_decomposition.py b/tests/link/mlx/linalg/test_decomposition.py index 6596fad627..85bb632c55 100644 --- a/tests/link/mlx/linalg/test_decomposition.py +++ b/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) diff --git a/tests/link/mlx/linalg/test_solvers.py b/tests/link/mlx/linalg/test_solvers.py index 0028296450..73444a8fc3 100644 --- a/tests/link/mlx/linalg/test_solvers.py +++ b/tests/link/mlx/linalg/test_solvers.py @@ -44,8 +44,14 @@ def test_mlx_solve(assume_a): ) +@pytest.mark.parametrize( + "unit_diagonal", [False, True], ids=["stored_diagonal", "unit_diagonal"] +) @pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"]) -def test_mlx_SolveTriangular(lower): +def test_mlx_SolveTriangular(lower, unit_diagonal): + # `unit_diagonal=True` was dropped on the way to `mx.linalg.solve_triangular`, + # which has no such argument, so the stored diagonal was used instead of ones + # and a wrong answer was returned without raising (#2384). rng = np.random.default_rng(15) A = pt.tensor("A", shape=(5, 5)) @@ -59,7 +65,7 @@ def test_mlx_SolveTriangular(lower): b, trans=0, lower=lower, - unit_diagonal=False, + unit_diagonal=unit_diagonal, ) compare_mlx_and_py( [A, b], @@ -128,3 +134,150 @@ def test_mlx_CholeskySolve_mixed_dtypes(): np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True ), ) + + +def test_mlx_SolveTriangular_unit_diagonal_lu_solve(): + # Batched `lu_factor` is blocked on #2385, so this covers the core case only. + batch_shape = () + # `lu_solve` packs a unit-triangular `L` into the LU array, with `U`'s + # diagonal occupying those slots, so it relies on `unit_diagonal=True` + # actually being honoured (#2384). + rng = np.random.default_rng(15) + n = 4 + + A = pt.tensor("A", shape=(*batch_shape, n, n)) + b = pt.tensor("b", shape=(*batch_shape, n)) + + A_val = rng.normal(size=(*batch_shape, 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=(*batch_shape, n)).astype(config.floatX) + + out = pt.linalg.lu_solve(pt.linalg.lu_factor(A), b, b_ndim=1) + + compare_mlx_and_py( + [A, b], + [out], + [A_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True + ), + ) + + +@pytest.mark.parametrize("batch_shape", [(4,), (2, 3)], ids=["batch", "batch_2d"]) +@pytest.mark.parametrize("b_ndim", [1, 2], ids=["b_vec", "b_mat"]) +def test_mlx_solve_batched(batch_shape, b_ndim): + """`Blockwise` used to `vmap` `mx.linalg.solve`, which has no `LUF` rule (#2385).""" + rng = np.random.default_rng(15) + n = 3 + b_shape = (*batch_shape, n) if b_ndim == 1 else (*batch_shape, n, 2) + + A = pt.tensor("A", shape=(*batch_shape, n, n)) + b = pt.tensor("b", shape=b_shape) + + A_val = rng.normal(size=(*batch_shape, 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_shape).astype(config.floatX) + + compare_mlx_and_py( + [A, b], + [pt.linalg.solve(A, b, b_ndim=b_ndim)], + [A_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True + ), + ) + + +@pytest.mark.parametrize( + "A_batch, b_batch", + [((4,), ()), ((), (4,)), ((1,), (4,))], + ids=["A_batched", "b_batched", "A_broadcast"], +) +def test_mlx_solve_batch_broadcast(A_batch, b_batch): + """`mx.linalg.solve` does not broadcast batch dims, so they must be aligned.""" + 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)) + + 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], + [pt.linalg.solve(A, b, b_ndim=1)], + [A_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True + ), + ) + + +@pytest.mark.parametrize( + "unit_diagonal", [False, True], ids=["stored_diagonal", "unit_diagonal"] +) +@pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"]) +def test_mlx_SolveTriangular_batched_ignores_other_triangle(lower, unit_diagonal): + """A batched triangular solve must still ignore the opposite triangle. + + `mx.vmap`'s `solve_triangular` rule drops the `upper` flag, so a batched + solve silently solved the *full* system. Nothing raised, and the answer was + only correct when the input happened to be exactly triangular already -- + which `cho_solve` and `lu_solve` never are, since they pack `L` and `U` + into a single array (#2385). + """ + rng = np.random.default_rng(15) + n, batch = 3, 4 + + A = pt.tensor("A", shape=(batch, n, n)) + b = pt.tensor("b", shape=(batch, n)) + + # Deliberately dense: both triangles are populated, so solving the full + # system gives a different answer from the triangular one. + A_val = rng.normal(size=(batch, n, n)).astype(config.floatX) + A_val += n * np.eye(n, dtype=config.floatX) + b_val = rng.normal(size=(batch, n)).astype(config.floatX) + + out = pt.linalg.solve_triangular( + A, b, lower=lower, unit_diagonal=unit_diagonal, b_ndim=1 + ) + compare_mlx_and_py( + [A, b], + [out], + [A_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True + ), + ) + + +@pytest.mark.parametrize("batch_shape", [(4,), (2, 3)], ids=["batch", "batch_2d"]) +def test_mlx_lu_solve_batched(batch_shape): + """Batched `lu_factor` + `lu_solve`: `LUF` vmap, the packed unit `L`, and + the pivot-to-permutation scan all had to be batched (#2385).""" + rng = np.random.default_rng(15) + n = 4 + + A = pt.tensor("A", shape=(*batch_shape, n, n)) + b = pt.tensor("b", shape=(*batch_shape, n)) + + A_val = rng.normal(size=(*batch_shape, 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=(*batch_shape, n)).astype(config.floatX) + + compare_mlx_and_py( + [A, b], + [pt.linalg.lu_solve(pt.linalg.lu_factor(A), b, b_ndim=1)], + [A_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-4, rtol=1e-4, strict=True + ), + ) diff --git a/tests/link/mlx/linalg/test_summary.py b/tests/link/mlx/linalg/test_summary.py index de587bf59b..2590ef0194 100644 --- a/tests/link/mlx/linalg/test_summary.py +++ b/tests/link/mlx/linalg/test_summary.py @@ -30,3 +30,15 @@ def test_mlx_slogdet(): sign, logabsdet = pt.linalg.slogdet(A) compare_mlx_and_py([A], [sign, logabsdet], [A_val], mlx_mode=get_mode("MLX")) + + +@pytest.mark.parametrize("batch_shape", [(4,), (2, 3)], ids=["batch", "batch_2d"]) +def test_mlx_det_batched(batch_shape): + """`Det` factorizes with `mx.linalg.lu_factor`, which `mx.vmap` cannot batch (#2385).""" + rng = np.random.default_rng(15) + n = 3 + + 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.det(A)], [A_val])