diff --git a/pytensor/link/mlx/dispatch/basic.py b/pytensor/link/mlx/dispatch/basic.py index 7c5f1e23a6..30ba1131e0 100644 --- a/pytensor/link/mlx/dispatch/basic.py +++ b/pytensor/link/mlx/dispatch/basic.py @@ -132,7 +132,9 @@ def mlx_typify_tensor(data, dtype=None, **kwargs): # and it does so on the CPU too, where float64 is perfectly usable if dtype is None and data.dtype == np.float64 and float64_supported(): dtype = mx.float64 - return _nan_safe_constant(data, dtype=dtype) + # MLX's elementwise kernels misread a non-contiguous buffer, and rewriting + # produces such arrays -- `triu` of a transpose, for one + return _nan_safe_constant(np.asarray(data, order="C"), dtype=dtype) @mlx_typify.register(slice) diff --git a/pytensor/link/mlx/dispatch/linalg/solvers.py b/pytensor/link/mlx/dispatch/linalg/solvers.py index 2f3f31416a..67c35b20fc 100644 --- a/pytensor/link/mlx/dispatch/linalg/solvers.py +++ b/pytensor/link/mlx/dispatch/linalg/solvers.py @@ -33,12 +33,22 @@ def solve(a, b): @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): + A = A.astype(stream=mx.cpu, dtype=A_dtype) + + if unit_diagonal: + # MLX's `solve_triangular` has no `unit_diagonal`. LAPACK's `trtrs` + # never reads the diagonal in that mode, so overwriting it with ones + # gives the same answer. + 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( - A.astype(stream=mx.cpu, dtype=A_dtype), + A, b.astype(stream=mx.cpu, dtype=b_dtype), upper=not lower, stream=mx.cpu, diff --git a/tests/link/mlx/linalg/test_decomposition.py b/tests/link/mlx/linalg/test_decomposition.py index 6596fad627..81295b6e68 100644 --- a/tests/link/mlx/linalg/test_decomposition.py +++ b/tests/link/mlx/linalg/test_decomposition.py @@ -5,7 +5,7 @@ from packaging.version import parse as V import pytensor.tensor as pt -from pytensor import config +from pytensor import config, grad from pytensor.tensor.linalg.decomposition import lu, svd from pytensor.tensor.linalg.decomposition.cholesky import cholesky from tests.link.mlx.test_basic import compare_mlx_and_py, mlx_mode @@ -131,6 +131,27 @@ def test_mlx_lu_factor(): compare_mlx_and_py([A], out, [A_val]) +def test_mlx_lu_solve_grad(): + """The pullback folds non-contiguous constants, which MLX must read correctly.""" + rng = np.random.default_rng(15) + + A = pt.tensor(name="A", shape=(5, 5)) + b = pt.tensor(name="b", shape=(5,)) + A_val = rng.normal(size=(5, 5)).astype(config.floatX) + b_val = rng.normal(size=(5,)).astype(config.floatX) + + x = pt.linalg.lu_solve(pt.linalg.lu_factor(A), b) + + # `fast_run` so the permutation `arange` in the pullback folds to a constant, + # which MLX's `arange` requires + compare_mlx_and_py( + [A, b], + [grad(x.sum(), A)], + [A_val, b_val], + mlx_mode=mlx_mode.including("fast_run"), + ) + + def test_mlx_pivot_to_permutations(): rng = np.random.default_rng(15) diff --git a/tests/link/mlx/linalg/test_solvers.py b/tests/link/mlx/linalg/test_solvers.py index 0028296450..1a4e88d9cc 100644 --- a/tests/link/mlx/linalg/test_solvers.py +++ b/tests/link/mlx/linalg/test_solvers.py @@ -44,14 +44,19 @@ def test_mlx_solve(assume_a): ) +@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): +def test_mlx_SolveTriangular(lower, unit_diagonal): rng = np.random.default_rng(15) A = pt.tensor("A", shape=(5, 5)) b = pt.tensor("B", 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) out = pt.linalg.solve_triangular( @@ -59,7 +64,7 @@ def test_mlx_SolveTriangular(lower): b, trans=0, lower=lower, - unit_diagonal=False, + unit_diagonal=unit_diagonal, ) compare_mlx_and_py( [A, b], diff --git a/tests/link/mlx/test_basic.py b/tests/link/mlx/test_basic.py index 0aa55c54eb..cdc1b23c53 100644 --- a/tests/link/mlx/test_basic.py +++ b/tests/link/mlx/test_basic.py @@ -23,7 +23,7 @@ mx = pytest.importorskip("mlx.core") -from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx +from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx, mlx_typify optimizer = RewriteDatabaseQuery(include=["mlx"], exclude=MLX._optimizer.exclude) @@ -355,3 +355,25 @@ def test_nan_array_constant(): compare_mlx_and_py( [x], [x + c], [np.array([10.0, 20.0, 30.0], dtype=config.floatX)] ) + + +@pytest.mark.parametrize( + "data", + [ + np.asfortranarray(np.arange(6, dtype="float32").reshape(2, 3)), + np.arange(6, dtype="float32").reshape(2, 3).T, + np.arange(6, dtype="float32").reshape(2, 3), + np.array(2.0, dtype="float32"), + ], + ids=["fortran_order", "transposed_view", "already_contiguous", "scalar"], +) +def test_mlx_typify_row_contiguous(data): + # MLX's elementwise kernels misread a non-contiguous buffer, and rewriting + # produces such arrays (``triu`` of a transpose, for one). Only the full LU + # pullback reproduces the wrong answer, so the invariant is pinned here. + result = mlx_typify(data) + as_numpy = np.asarray(result) + + np.testing.assert_array_equal(as_numpy, data) + assert as_numpy.flags["C_CONTIGUOUS"] + assert result.shape == data.shape