From 5e363328b11cae75dcc6d940b9e529061c16bf50 Mon Sep 17 00:00:00 2001 From: guillaume-osmo Date: Wed, 26 Aug 2026 15:51:12 +0200 Subject: [PATCH] Honour unit_diagonal in the MLX SolveTriangular dispatch (#2384) `mlx_funcify_SolveTriangular` read `op.lower` but never `op.unit_diagonal`, and `mx.linalg.solve_triangular` has no such argument, so the flag was dropped and the stored diagonal used instead of ones. Nothing raised. This also silently broke `lu_solve`, whose unit-triangular `L` is packed into the LU array with `U`'s diagonal in those slots, so any graph where `reuse_decomposition_multiple_solves` collapses two solves onto one factorization returned wrong numbers on MLX. Overwrite the diagonal with ones before the solve when the flag is set; this broadcasts over leading batch dims, so the batched case is covered too. `test_mlx_SolveTriangular` now parametrizes over `unit_diagonal`, and a new test pins the `lu_factor` + `lu_solve` path. --- pytensor/link/mlx/dispatch/linalg/solvers.py | 25 ++++++++++--- tests/link/mlx/linalg/test_solvers.py | 39 +++++++++++++++++++- 2 files changed, 56 insertions(+), 8 deletions(-) diff --git a/pytensor/link/mlx/dispatch/linalg/solvers.py b/pytensor/link/mlx/dispatch/linalg/solvers.py index 2f3f31416a..d021979950 100644 --- a/pytensor/link/mlx/dispatch/linalg/solvers.py +++ b/pytensor/link/mlx/dispatch/linalg/solvers.py @@ -30,19 +30,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 diff --git a/tests/link/mlx/linalg/test_solvers.py b/tests/link/mlx/linalg/test_solvers.py index 0028296450..5a248e6d3d 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,32 @@ 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 + ), + )