linalg: solve(A, b) via the on-ANE pivoted LU, and a matrix right-hand side - #170
Merged
sbryngelson merged 1 commit intoAug 4, 2026
Merged
Conversation
…d side Implements sbryngelson#132. `lu_pivoted` already factors on the ANE and `solve_triangular` already substitutes, so `solve` is the permutation applied host-side plus two substitutions. It is the direct counterpart to `conjugate_gradient`/`gmres`, which are iterative and want SPD input, and it handles general square A. `solve_triangular` now takes a matrix right-hand side as well as a vector, which sbryngelson#132 asks for and sbryngelson#106 wants for `inv`. The substitution is a row recurrence, so widening the row slice from [1,1] to [1,m] solves all m columns in the same fused graph rather than dispatching per column. A 1-D b still returns shape [n]. Rejects a non-square A, a b whose rows do not match, and a singular factor (zero pivot in U) with LinAlgError, like numpy. On M2 Pro against numpy, relerr <= 1.3e-3 for n in {4,6,8} for both vector and matrix right-hand sides; the matrix path also agrees with solving each column on its own, and a residual check confirms A x ~ b without a reference solver. test_linalg.py 66 passed, corpus GATE GREEN 90/90.
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
2 participants
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.
Fixes #132.
What lands
solve(A, b)for general square A:lu_pivotedalready factors on the ANE andsolve_triangularalready substitutes, so this is the permutation applied host-side (P is a permutation matrix, so P b is a row gather rather than a matmul worth dispatching) plus forward and back substitution. It is the direct complement toconjugate_gradient/gmres, which are iterative and want SPD input.solve_triangularnow accepts a matrix right-hand side, which #132 asks for and which #106 wants forinv. The substitution is a row recurrence, so widening the row slice from[1, 1]to[1, m]solves all m columns in the same fused graph instead of dispatching once per column. A 1-Dbstill returns shape[n], unchanged.Both are added to
__all__, along withsolve_triangular, which was reachable but not exported.Verification
Apple M2 Pro (Mac14,12 Mac mini), macOS 26.5.2. Against numpy on the fp16-rounded system:
tests/test_linalg.pytests/run_corpus.pyruff/pyright/.githooks/pre-commitThe new tests fail without the change, verified by swapping in
origin/main'slinalg.py.Three of them are worth pointing at specifically, since "matches numpy" alone would be weak here:
A x ~ b) confirms correctness without a reference solver at alltest_solve_triangular_vector_shape_unchangedpins that 1-D input still returns 1-D, which is the regression the widening could most easily have introducedRejects covered: non-square A, mismatched b rows, and a singular factor raising
LinAlgErrorlike numpy.Risk
Contained to
linalg.py. The change tosolve_triangularis the one with existing callers, so the vector shape and the existing tests (includingtest_solve_triangular_large_entries, which exercises the routed accessor above the slice saturation threshold) are the thing to look at; both are unchanged and green.expm(#133) is the natural next one and I have it working locally, but it is a separate issue so it goes in its own PR rather than riding along here.