refactor: create invert tags helper (#224)#5
Open
jpbrodrick89 wants to merge 16 commits into
Open
Conversation
Add a Hermitian eigenvalue decomposition solver (`lineax.HEVD`) built on `jax.numpy.linalg.eigh`, mirroring `lineax.SVD` for self-adjoint operators. Like SVD it handles singular operators by returning the pseudoinverse solution, and honours `MaxRankTag` by statically truncating to the largest eigenvalues. Unlike SVD's descending singular values, `eigh` returns eigenvalues in ascending order and they are signed, so the decomposition is reordered by descending magnitude before the (shared) rcond/truncation logic. Add a `hermitian_tag` and `is_hermitian` check mirroring the symmetric ones, including dispatch for all built-in operators, tag propagation through transpose/inversion/composition/scaling, and the real-vs-complex relationships between symmetry and Hermitian-ness.
…rmitian `compute` masks small eigenvalues by magnitude regardless of ordering, so the argsort + column gather (O(n^2)) is only needed to enable the static low-rank truncation. Move it inside the `MaxRankTag` branch so the common path keeps `eigh`'s native order at no cost, and have `compute` take `max|w|` directly. Dispatch `AutoLinearSolver(well_posed=False)` to `HEVD` for Hermitian operators (a Hermitian eigendecomposition is cheaper than a general SVD, with the same pseudoinverse handling of ill-posed systems).
Add tests building `Q diag(eigvals) Q^H` with zero eigenvalues placed between large eigenvalues of both signs. After eigh's ascending sort the discarded eigenpairs sit in the interior of the spectrum, exercising the magnitude-based masking (pseudoinverse) and the magnitude-based MaxRankTag truncation, for both real and complex dtypes.
`AbstractLinearOperator.H` returns the Hermitian adjoint: a no-op (`self`) for a Hermitian operator (since `Aᴴ = A`), otherwise `conj(operator).transpose()`. This mirrors the existing `.T`, whose `transpose()` already no-ops for symmetric operators. Add `AbstractLinearSolver.conj_transpose(state, options)` with a default that composes `conj` then `transpose` (identical to the previous inline computation in the linear-solve JVP). Self-adjoint solvers (`Cholesky`, `HEVD`) override it as a no-op since `Aᴴ = A` -- note their `transpose` must NOT no-op, as `Aᵀ = conj(A) != A` for complex operators. `AutoLinearSolver` delegates so the selected solver's override applies. Wire both into the JVP rule: `operator.H` / `t_operator.H` and `solver.conj_transpose(state)` replace the explicit conj+transpose, skipping redundant work on the common self-adjoint paths.
Overriding the concrete `conj_transpose` in Cholesky/HEVD/AutoLinearSolver violated equinox's abstract/final convention. Instead make `conj_transpose` a single final method on `AbstractLinearSolver` that no-ops when a new abstract `assume_hermitian()` returns True (else composes `conj` then `transpose`), and implement `assume_hermitian` on every solver -- mirroring the existing `assume_full_rank`. Cholesky and HEVD return True; all others return False (`AutoLinearSolver` conservatively False, since its solver is chosen at runtime).
…_transpose Whether the conjugate-transpose state can be reused is fundamentally an operator property: for a Hermitian operator `Aᴴ = A`, so `init(Aᴴ) == init(A)` for *any* solver. So the JVP now skips forming the adjoint state when `is_hermitian(operator)` and reuses the existing state directly -- a single operator-keyed check that applies uniformly (including `AutoLinearSolver`, CG, etc.), rather than a per-solver flag. This removes the need for both the `assume_hermitian` abstract method (reverted from every solver) and the `conj_transpose` solver method (its only caller was the JVP, which now inlines `conj` then `transpose` for the non-Hermitian case, exactly as before). The operator-level `.H` no-op is retained.
A PSD/NSD operator is Hermitian (real or complex), and a real symmetric/diagonal operator is Hermitian too. The generic tag loop for TaggedLinearOperator only checked `hermitian_tag in tags or is_hermitian(inner)`, so a `TaggedLinearOperator` carrying e.g. `positive_semidefinite_tag` (as produced by composing operators) reported `is_hermitian == False`, causing HEVD to reject it. Give TaggedLinearOperator an explicit `is_hermitian` mirroring the cross-implications encoded for the core operators. Fixes the test_singular failures for HEVD on make_composed_operator (PSD/NSD, real and complex).
HEVD is not special enough to warrant a bespoke test module. Instead: - Make the generic `zero` singular construction yield a genuinely rank-deficient *indefinite* Hermitian matrix (re-zero the leading row and column after the symmetric/Hermitian build), so HEVD with `hermitian_tag` can join the standard pseudoinverse scan (forward solve, transpose, linearise, materialise, jvp, vmap, real and complex). This removes the need for bespoke indefinite/singular tests. - Move the MaxRankTag truncation and over-truncation tests to test_max_rank.py, alongside the SVD equivalents. - Move the is_hermitian checks and hermitian-tag propagation to test_operator.py, alongside is_symmetric. - Move the `.H` (conjugate-transpose) and adjoint-state-reuse tests to test_transpose.py. - Drop the redundant solve/pseudoinverse/jit/grad/rejection tests (covered by the generic scans) and delete test_hevd.py.
For a PSD (resp. NSD) operator all eigenvalues are >= 0 (resp. <= 0), so in eigh's ascending order the largest-magnitude `r` are a contiguous trailing (resp. leading) slice. Querying the PSD/NSD tags lets MaxRankTag truncation take that slice directly, avoiding the argsort + column gather that the indefinite case still needs (slices are much cheaper than gathers, especially on accelerators). The branching is entirely in `init` (static, at trace time); `compute` is unchanged. The over-truncation check now tests the largest *discarded* magnitude rather than the boundary eigenvalue, which also guards against a mistagged PSD/NSD operator whose true large eigenvalues sit on the dropped side. Parametrise the max_rank tests over all three branches.
Every non-well-posed solver has a "gram partner": a solver representing the (pseudo)inverse of the gram matrix AᴴA, obtainable from the existing factorisation with no further decomposition (QR -> Cholesky with R; SVD/HEVD -> HEVD with σ²/w²; Normal(tall) -> its inner solver, which already factorises AᴴA). The `not assume_independent_rows` term of the JVP is `A⁺(Aᴴ)⁺w = (AᴴA)⁺w`. Use the gram partner to compute this in a single gram solve instead of nesting an adjoint solve inside the outer A⁺ solve. Routing it back through `linear_solve_p` (rather than applying the factors directly) keeps it correct under higher-order autodiff, since the gram solve then uses lineax's pseudoinverse-aware adjoint rather than differentiating through the factorisation. Contained in `_solve.py` via a `_gram_partner` dispatch (lazy imports to avoid the solver<->_solve cycle; `_HasGramPartner` type alias for readability). The row-space projection term is deliberately left untouched.
The gram operator AᴴA depends only on `operator`, so construct it inline at the JVP call site (where `operator` is in scope) and drop the `_gram_operator` helper and the operator from `_gram_partner`'s return (now `(gram_solver, gram_state)`); the SVD/HEVD branch takes it as an argument for `pack_structures`. QR only reaches the gram fast path for tall operators (`not assume_independent_rows` requires `rows > columns`), where it stored the QR of `A`, so the wide/transposed case is unreachable -- raise instead of returning NotImplemented. (Normal's wide case, by contrast, is genuinely reachable for a rank-deficient inner solver, so it keeps the NotImplemented fallback.)
…r' into claude/hevd-solver-hermitian-1paoqv # Conflicts: # lineax/_solve.py
jpbrodrick89
changed the base branch from
dev
to
jpb/refactor-AbstractLinearSolver
June 18, 2026 20:57
…cate Now that the AbstractLinearSolver refactor lets `_solve.py` import the solver classes at module scope, gate the gram fast path on a `_has_gram_partner(solver, state)` predicate instead of the `NotImplemented` dance: - `_MaybeHasGramPartner = QR | SVD | HEVD | Normal` names the candidate solver *types* (renamed from `_HasGramPartner`, since Normal only conditionally has one). - `_has_gram_partner(solver, state)` is the definitive runtime check (True for QR/SVD/HEVD; for Normal, only when tall). - `_gram_partner` is typed `_MaybeHasGramPartner -> tuple[...]`, always returns a pair, and raises ValueError on the wide cases as guards (unreachable given the predicate). Also only construct the (lazy) gram operator when the fast path is actually taken.
Unlike the other property checks, `is_hermitian` is an optional optimisation hint (dispatch to HEVD; shortcut the linear-solve JVP adjoint state), and `False` is always the safe answer. Defaulting to `False` rather than `_default_not_implemented` means a custom `AbstractLinearOperator` written before `is_hermitian` existed keeps working without registering it -- otherwise it would raise in `__check_init__`, the JVP, and `AutoLinearSolver`. Built-in operators all register `is_hermitian`, so they are unaffected.
A plain `return False` default was wrong for custom operators: a PSD custom operator would report `is_hermitian == False`, inconsistent with built-in PSD operators (which are Hermitian) -- so it would miss HEVD dispatch and the JVP shortcut that an equivalent built-in gets. Instead default to the same logic the built-in operators use (minus the `hermitian_tag` check, which needs operator-specific `.tags`): Hermitian if PSD/NSD, or if real and symmetric. This defers to checks a complete custom operator already registers (and that `__check_init__`/`Cholesky` already require), so it stays non-breaking while keeping `is_hermitian` correct -- the graceful-default pattern also used by `max_rank`.
It is now used across several files in the `_operator` package (the symmetric/Hermitian checks), so drop the leading underscore to mark it as a shared package helper.
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.
No description provided.