Honour unit_diagonal in the MLX SolveTriangular dispatch - #2396
Open
guillaume-osmo wants to merge 1 commit into
Open
Honour unit_diagonal in the MLX SolveTriangular dispatch#2396guillaume-osmo wants to merge 1 commit into
guillaume-osmo wants to merge 1 commit into
Conversation
) `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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #2384.
Motivation
mlx_funcify_SolveTriangularreadop.lowerbut neverop.unit_diagonal, andmx.linalg.solve_triangularhas no such argument, so the flag was dropped and the stored diagonal used instead of ones. Nothing raised.As the issue notes, this also silently broke
lu_solve, whose unit-triangularLis packed into the LU array withU's diagonal occupying those slots — so any graph wherereuse_decomposition_multiple_solvescollapses two solves onto one factorization returned wrong numbers.Implementation
Overwrite the diagonal with ones before the solve when the flag is set. The helper broadcasts over leading batch dims, so the batched case is covered.
Tests
test_mlx_SolveTriangularnow parametrizes overunit_diagonal(× the existinglower), so the four combinations are covered.test_mlx_SolveTriangular_unit_diagonal_lu_solvepins thelu_factor+lu_solvepath the issue calls out.The batched
lu_solvecase is deliberately not covered here — it is blocked on #2385, and is added in that PR instead.Full
tests/link/mlx/suite: no regressions.