Add determinant support#6
Open
jpbrodrick89 wants to merge 34 commits into
Open
Conversation
…determinant functions Introduces AbstractDirectLinearSolver as an intermediate abstract class between AbstractLinearSolver and the concrete direct solvers (LU, Cholesky, QR, SVD, Triangular, Diagonal, Tridiagonal). Direct solvers materialise the operator and can therefore expose the (log absolute) determinant from their factored state. Abstract methods: - logabsdet(state, options) -> log|det(A)| - det_sign(state, options) -> sign(det(A)) Concrete implementations: - LU: sum(log|diag(U)|), (-1)^swaps * prod(sign(diag(U))) - Cholesky: 2*sum(log|diag(factor)|), +/-1 based on PSD/NSD and matrix size - QR: raises NotImplementedError (geqrf has no JAX AD rules; deferred follow-up) - SVD: sum(log(s)) for square operators; det_sign raises (O(n³) to compute) - Triangular: sum(log|diag|), prod(sign(diag)); short-circuits for unit diagonal - Diagonal: sum(log|diag|), prod(sign(diag)); short-circuits for unit diagonal - Tridiagonal: Thomas algorithm pivot recurrence (O(n), AD-friendly via lax.scan) Normal is refactored with a __new__ factory that returns _NormalDirect when the inner solver is an AbstractDirectLinearSolver. _NormalDirect inherits all Normal logic and adds logabsdet via: 0.5 * inner_solver.logabsdet(inner_state, ...). det_sign is not supported (squaring destroys sign information). The Normal.init materialise check is generalised from isinstance(_, Cholesky) to isinstance(_, AbstractDirectLinearSolver). Top-level convenience functions added to lx namespace: - lx.logabsdet(operator, solver) -> log|det| - lx.slogdet(operator, solver) -> (sign, log|det|) [numpy convention] - lx.determinant(operator, solver) -> det https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Reverts the _NormalDirect subclass approach which violated equinox's abstract-or-final pattern (Normal was both instantiatable and subclassable). Instead, Normal.logabsdet and Normal.det_sign are added as plain concrete methods directly on Normal (no inheritance from AbstractDirectLinearSolver). logabsdet delegates to inner_solver.logabsdet when the inner solver is direct, and raises TypeError otherwise. det_sign always raises NotImplementedError. The is_direct(solver) helper is added to normal.py and exported at the top level. It returns True for AbstractDirectLinearSolver instances and for Normal instances whose inner_solver is direct. Normal.init uses is_direct instead of the previous isinstance(_, AbstractDirectLinearSolver) inline check. The top-level lx.logabsdet/slogdet/determinant functions accept AbstractDirectLinearSolver as their type hint; Normal works via duck typing at runtime when wrapping a direct inner solver. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Normal(Normal(LU())) is contrived but is_direct should handle it correctly. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
…ints Replaces the AbstractDirectLinearSolver type hint on lx.logabsdet, lx.slogdet, and lx.determinant with a SupportsLogdet Protocol. This avoids the circular import that would prevent a union type (AbstractDirectLinearSolver | Normal) from living in _solve.py, and is forward-compatible with any third-party solver that implements the three required methods (init, logabsdet, det_sign). SupportsLogdet is exported at the top level. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
…irectLinearSolver interface Protocol cannot inherit from AbstractLinearSolver due to equinox/typing metaclass conflict, so all seven methods are listed explicitly to stay consistent with AbstractDirectLinearSolver's interface. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
…ctions to _determinant.py Uses union type AbstractDirectLinearSolver | Normal instead of a Protocol, avoiding metaclass conflicts and making accepted types explicit in docs. Circular import avoidance: _solve.py cannot import Normal (which imports from _solve.py), so _determinant.py imports from both. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
…ndefined Sign is returned as nan (not raised) for Normal (gram matrix destroys sign information) and SVD (recovering sign(det(U)) and sign(det(V^T)) requires O(n^3) extra work). QR still raises NotImplementedError since geqrf has no JAX AD rules, which is a categorically different limitation. lx.determinant gains a throw=True kwarg (matching linear_solve convention): eqx.error_if fires on isnan(sign) with a branching message by solver type, operator shape (non-square), and assume_full_rank. Bonus: Tridiagonal.slogdet computes the pivot scan once instead of twice. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
…rove SVD sign API changes: - Remove lx.logabsdet (users can inline `_, lad = lx.slogdet(...)`) - Remove MaybeDirectLinearSolver type alias; inline AbstractDirectLinearSolver | Normal in function signatures for clearer documentation slogdet/determinant: accept optional state= kwarg (no stop_gradient, unlike linear_solve); docstring warns callers not to stop_gradient the state since determinants are computed directly from the factorisation. SVD.slogdet: detect rank deficiency via rcond threshold and return (0, -inf) for rank-deficient matrices (det = 0 by definition, sign = 0 trivially). For full-rank square matrices sign is still nan (requires O(n^3) extra work). This avoids a spurious throw=True error when the determinant is simply zero. Error messages: dispatch assume_full_rank first (same message whether or not Normal is involved), then Normal-specific message ending with "use lx.QR() for the pseudodeterminant of a full-rank rectangular matrix". All messages written assuming QR sign support will be available before merge. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
SVD operates in the pseudo paradigm throughout:
- lad = sum(log(non-zero singular values)) for all shapes/ranks, including
rank-deficient operators. The zero singular values are excluded via the
rcond threshold (replaced with 1 so log = 0, they don't contribute).
- sign = nan always — SVD cannot recover the sign in any case:
full-rank square: needs sign(det(U)) * sign(det(V^T)), O(n^3)
full-rank rectangular: future work via QR Householder vectors
rank-deficient: needs an eigensolver for the pseudodeterminant sign
Previously returned sign=0 for rank-deficient matrices (treating it as
a true determinant rather than a pseudodeterminant), which is wrong under
the pseudodeterminant convention: the pseudodeterminant of a non-zero
rank-deficient matrix is non-zero. Also removed the ValueError for
non-square operators — rectangular is a perfectly valid input.
https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Mirrors linear_solve's design: stop_gradient the state and provide an
analytic derivative rather than differentiating through the factorisation.
JVP rule: d(log|det(A)|)/dA = trace(A† dA), computed by vmapping
linear_solve over the columns of dA (the tangent operator materialised
via TangentLinearOperator.as_matrix()). Each column solve reuses the
stopped-gradient state from the primal for efficiency.
For complex operators: the imaginary part of trace(A^{-1} dA) is
transferred to sign_dot (matching JAX's own slogdet convention);
the real part becomes lad_dot.
determinant inherits the correct JVP automatically:
d(det(A))/dA = det(A) * trace(A^{-1} dA) * A^{-T}
which equals det(A) * A^{-T} as expected.
https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
QR assumes full rank, so sign(det) = sign(det(Q)) * sign(det(R)) is well-defined for tall, wide, and square operators alike. Remove the is_square guard that was returning nan for non-square inputs. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
…tion slogdet implementations in qr/svd/cholesky/normal used jnp.result_type to compute float_dtype, which raises TypePromotionError under strict dtype promotion (used in tests). Switch to numpy.result_type, which ignores JAX's promotion config. Tests cover: square determinant correctness, sign=nan for SVD/Normal, throw kwarg, QR rectangular (lad vs SVD, sign=±1, explicit QR reference), state= kwarg reuse, and JVP of slogdet/determinant. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
- Drop unnecessary float32 floor in slogdet dtype computation across qr/svd/cholesky/normal; solvers always operate on float types so just use .real.dtype directly - Add Normal(Cholesky) lad correctness test against jnp.linalg.slogdet - JVP tests now compare against jax.jvp of jnp.linalg.slogdet (not just the analytical formula), and cover use_state=True - Add grad (VJP) test comparing against jax.grad of jnp.linalg.slogdet - Drop trivial state=kwarg standalone test; state is exercised in JVP https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Normal is designed for non-square operators. Test with tall (5×3) and wide (3×5) matrices; validate lad against sum(log(singular values)). https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
- Rename _slogdet_p -> _slogdet (not a JAX primitive) - Move sentinel dance and stop_gradient into the public slogdet/determinant functions, matching linear_solve's structure - Default state=sentinel instead of state=None - Mark state as nondifferentiable inside _slogdet via eqxi.nondifferentiable - Factor out _prep_state helper shared by slogdet and determinant https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
_prep_state was only called in one place after determinant was changed to call slogdet directly. Inline the sentinel/stop_gradient logic into slogdet and remove the helper. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
- eqxi.nondifferentiable(state) now lives in the public slogdet function (after stop_gradient), not inside _slogdet, matching linear_solve's pattern of preparing all inputs before the inner call - Add IdentityLinearOperator fast path before the sentinel dance: det(I) = 1, so return (sign=1, lad=0) without any factorisation https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
- Parametrize square correctness and AD tests over make_matrix_operator and make_jac_operator, matching test_jvp.py's pattern; JacobianLinearOperator exercises the TangentLinearOperator.as_matrix() path for JVP/grad - Add test_slogdet_jvp_jvp: second-order JVP compared against jnp.linalg.slogdet for LU and QR https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
When well_posed=False, zero (or near-zero) diagonal entries are excluded from the determinant, using the same rcond threshold as compute(). This mirrors SVD's pseudodeterminant convention and avoids returning sign=0 or lad=-inf for rank-deficient diagonal operators. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
operator is the only differentiable argument to _slogdet (solver/options/state are all nondifferentiable), so t_operator is always present when the JVP is called. The has_t_op guard and _is_none helper were dead code inherited from linear_solve_p_jvp where both operator and vector can independently vary. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
- Drop the stop_gradient warning from slogdet/determinant state= docs; it's an implementation detail users don't need to know about - Normal.slogdet error message recommends lx.Cholesky() only, not LU — Normal wraps positive (semi)definite gram matrices so Cholesky is the natural pairing https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
AutoLinearSolver always delegates to direct solvers (LU, Cholesky, QR, SVD, Diagonal, Triangular, Tridiagonal), so it should satisfy the AbstractDirectLinearSolver interface. This lets it be used with lx.slogdet and lx.determinant directly. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Tests well_posed=True and well_posed=None; well_posed=False is excluded for the same reason SVD is excluded — sign is always nan from SVD. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
- New determinant.md page (mirrors linear_solve.md) documenting lx.slogdet and lx.determinant, including a note on XLA CSE - AbstractDirectLinearSolver added to solvers.md after AbstractLinearSolver, documenting its slogdet method - determinant.md added to mkdocs.yml nav after linear_solve.md https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Merge origin/jpb/hevd into the determinant branch, resolving conflicts: - AbstractLinearSolver and new AbstractDirectLinearSolver now live in _solver/base.py (moved from _solve.py during the operator refactor) - AutoLinearSolver moved to _solver/auto.py, promoted to AbstractDirectLinearSolver with slogdet delegation - HEVD promoted to AbstractDirectLinearSolver with slogdet: sign = prod(sign(w_i)) for non-masked eigenvalues (same rcond as compute()), lad = sum(log|w_i|) — full pseudodeterminant support including rank-deficient Hermitian matrices - Error message in determinant() updated to recommend HEVD for Hermitian/rank-deficient systems - Restore _check_rank_compat dropped during conflict resolution - HEVD added to SQUARE_DET_CASES test parametrisation https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
LU and QR assume full rank and won't work for rank-deficient systems. The message now recommends HEVD for Hermitian rank-deficient matrices and directs users to jnp.linalg.eig for non-Hermitian rank-deficient systems where lineax has no sign-capable solver. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
jpbrodrick89
changed the title
For complex matrices jnp.sign returns a unit complex number (the phase of the determinant), matching numpy.linalg.slogdet convention. The .real / .astype(real_dtype) casts were silently discarding that phase. Cholesky is unaffected (sign is always ±1 from a static Python int). HEVD is unaffected (Hermitian eigenvalues are always real). All other solvers (LU, QR, Diagonal, Triangular, Tridiagonal) now return the natural dtype of jnp.sign, preserving the complex phase. https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
The previous formula used prod(where(taus != 0, -1, 1)), which gives ±1 parity but is only correct for real Householder reflectors. For complex matrices, det(I - τ·v·vᴴ) = 1 - τ·‖v‖² (matrix determinant lemma), which yields a general unit complex number rather than always ±1. The new formula computes col_norms_sq from the subdiagonal of a (where the Householder vectors are stored in raw QR layout), then accumulates sign_Q = prod(1 - τ·‖v‖²). For real matrices, ‖v‖² = 1 + s² and τ = 2/(1+s²), so 1 - τ·‖v‖² = -1, matching the old behaviour exactly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
AbstractDirectLinearSolver is the public API for this concept. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ
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.
Summary
This PR refactors the monolithic
_operator.pymodule into a structured package with multiple submodules, while adding support for Hermitian operators and rank constraints viaMaxRankTag.Key Changes
Module Restructuring
lineax/_operator.py(2408 lines) into a packagelineax/_operator/:base.py: Abstract base class and utility functionscore.py: Core operators (MatrixLinearOperator,PyTreeLinearOperator,FunctionLinearOperator)structured.py: Structured operators (IdentityLinearOperator,DiagonalLinearOperator,TridiagonalLinearOperator)wrapper.py: Wrapper operators (TaggedLinearOperator,TangentLinearOperator, etc.)binary.py: Binary operators (AddLinearOperator,ComposedLinearOperator, etc.)__init__.py: Public API exportsNew Features
hermitian_tagandis_hermitian()function for complex matrices whereA = A^H(conjugate transpose)MaxRankTagdataclass to mark operators with bounded rank, queryable viamax_rank()function.Hproperty toAbstractLinearOperatoras shorthand for conjugate transposeSolver Module Reorganization
lineax/_solver/base.pylineax/_solver/auto.pyforAutoLinearSolverNew Determinant Module
lineax/_determinant.pywith determinant computation functionsdocs/api/determinant.mdTest Coverage
tests/test_max_rank.pywith comprehensiveMaxRankTagsemantics teststests/test_determinant.pywith determinant correctness testsDocumentation
is_hermitian,max_rank, and determinant functions.Hproperty to operator API documentationImplementation Details
MaxRankTagis a frozen dataclass with equality/hashing for use in frozensets@operator usesmin(rank_A, rank_B),+usesmin(rank_A + rank_B, in_size, out_size)__init__.pyre-exports_solve.pyhttps://claude.ai/code/session_01TywaHs2ctbDsGhrchbyXzJ