Add implicitly-differentiable Davidson eigensolvers (eigh + eig) - #119
Add implicitly-differentiable Davidson eigensolvers (eigh + eig)#119claude[bot] wants to merge 4 commits into
Conversation
Add pyscfad/lib/davidson_lite.py providing fully jittable Davidson eigensolvers differentiated via jax.lax.custom_root (implicit function theorem) instead of unrolling the iterations: - eigh: Hermitian/real-symmetric, supports complex Hermitian; degenerate response projected out (deg_thresh), mirroring _eigh_gen_jvp_rule. - eig: non-Hermitian, adjoint via explicit aopT to recover biorthonormal left vectors and oblique projectors. Converged states are recovered from the linearization rather than closed over, enabling correct higher-order derivatives. Forward solve is a fixed-max-space, GPU-friendly Davidson (while_loop, static shapes, batched aop, thick restart). Fixed nroots keeps it jittable. Add tests/test_davidson_lite.py: forward accuracy, jit, Hellmann-Feynman, check_grads(order=2, fwd+rev), eigenvector gradient vs dense eigh, degenerate spectrum, complex Hermitian, and non-Hermitian eig forward + biorthonormality + eigenvalue gradients. Co-authored-by: Xing Zhang <fishjojo@users.noreply.github.com>
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 71d6a6c728
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| restart = (nv + k) > m | ||
| Vb, off, valid = lax.cond( | ||
| restart, | ||
| lambda: (jnp.zeros_like(V).at[:, :k].set(X), jnp.int32(k), jnp.int32(k)), |
There was a problem hiding this comment.
When eig needs a thick restart (nv + k > max_space) with nroots > 1, this seeds the next basis with X directly, but the right Ritz vectors returned by jnp.linalg.eig are not mutually orthogonal. The later projected solve still assumes V.conj().T @ AV is built from an orthonormal basis, so after a restart it can select spurious near-zero Ritz pairs; for example, an 80x80 diagonally dominant matrix with nroots=2 and the default max_space returns eigenvalues around 0 instead of the lowest values near 1, 2, with tiny returned vectors satisfying the absolute residual test. Re-orthonormalizing X before storing it in V (and keeping the projected quantities consistent) avoids corrupting restarted non-Hermitian solves.
Useful? React with 👍 / 👎.
|
@claude address codex's review. |
|
Claude finished @fishjojo's task in 2m 57s —— View job Addressed Codex's review
The fix. Codex was right. On a thick restart ( I now re-orthonormalize the Ritz vectors via QR before storing them on restart: Xr = jnp.linalg.qr(X)[0]
Vb, off, valid = lax.cond(
restart,
lambda: (jnp.zeros_like(V).at[:, :k].set(Xr), jnp.int32(k), jnp.int32(k)),
lambda: (V, nv, nv))For the Hermitian ( Regression test. Added I verified the test fails without the fix (eigenvalue error ≈ 2.22, i.e. it returned roots near 0) and passes with it. Full file: 11/11 tests pass on jax 0.10.1, float64. |
On a thick restart in `eig`, the right Ritz vectors X = V @ C (with C from jnp.linalg.eig) are not mutually orthonormal, so seeding the next basis with them corrupted the projected solve V^H A V (which assumes an orthonormal basis) and could select spurious near-zero Ritz pairs. Re-orthonormalize X via QR before storing it in V. For the Hermitian case X is already orthonormal, so QR is a no-op up to column phase. Adds a regression test (80x80 diagonally dominant, nroots=2, default max_space) that forces multiple restarts. Co-authored-by: Xing Zhang <fishjojo@users.noreply.github.com>
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 374131e911
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| X, Y, w = state | ||
| R1 = aop(X) - X * w[None, :] | ||
| R2 = aopT(Y) - Y * jnp.conj(w)[None, :] | ||
| s = jnp.sum(jnp.conj(Y) * X, axis=0) - 1.0 |
There was a problem hiding this comment.
When differentiating any result that depends on the non-Hermitian eigenvectors, this constraint leaves the scaling gauge underdetermined: X -> a X, Y -> Y/conj(a) still satisfies the two eigen-equations and Y^H X = 1, while the forward solver returns unit-norm Ritz vectors from jnp.linalg.eig. The tangent solve therefore chooses a different gauge than the returned X; for example, for a small real diagonally dominant matrix, jax.grad(lambda p: sum((c*eig(...)[1][:,0].real)**2))(p) disagrees in sign/magnitude with finite differences/dense eigvec derivatives even though eigenvalue gradients pass. Please add constraints matching the forward normalization/phase (or avoid exposing differentiable X/Y for eig).
Useful? React with 👍 / 👎.
|
@claude address codex's latest review. |
|
Claude finished @fishjojo's task in 7m 38s —— View job Addressed Codex's latest review (P2: non-Hermitian eigenvector gauge)
The issue. Codex was right. For The fix. Took Codex's recommended alternative — avoid exposing differentiable X, Y, w = custom_root(root_fn, state0, solve, tangent_solve, has_aux=False)
# ... gauge note ...
X = lax.stop_gradient(X)
Y = lax.stop_gradient(Y)Important subtlety. I first tried short-circuiting Regression test. Added 12/12 tests pass on jax 0.10.1, float64. |
The right/left eigenvectors of a general operator have an undetermined complex scaling gauge that the forward jnp.linalg.eig normalization pins non-smoothly, so their derivatives cannot correspond to the returned vectors. Expose only the gauge-invariant eigenvalues as differentiable and stop gradients on the eigenvectors. The eigenvector tangents are still computed internally so the eigenvalue's higher-order derivatives remain correct. Co-authored-by: Xing Zhang <fishjojo@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 91cde4db9f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| X = lax.stop_gradient(X) | ||
| Y = lax.stop_gradient(Y) |
There was a problem hiding this comment.
When callers use return_left=True to differentiate a gauge-invariant quantity built from the biorthonormal pair, such as a spectral projector X @ Y.conj().T, this stop_gradient makes the entire eigenspace contribution differentiate as exactly zero even though the projector changes with the operator. I checked this with a simple diagonally dominant matrix: a projector scalar had finite-difference slope about -0.716 while jax.grad through eig(..., return_left=True) returned 0.0. The docstring explicitly points users at spectral projectors as the safe eigenvector-dependent quantity, so stopping both vectors here silently corrupts those gradients; stop only the non-differentiable gauge/phase or document that only eigenvalues may be differentiated.
Useful? React with 👍 / 👎.
|
@claude address codex's latest review. |
|
Claude finished @fishjojo's task in 4m 5s —— View job Addressed Codex's latest review (P2: projector gradients vs. stopped eigenvectors)
The issue. Codex correctly spotted a contradiction left by the previous round. The last fix wrapped both The fix — chose the conservative resolution (docstring + lock). Codex offered two mutually-exclusive options: (1) stop only the gauge/phase so projectors stay differentiable, or (2) document that only eigenvalues may be differentiated. I took option 2 because option 1 re-exposes differentiable
Regression test. Extended 12/12 tests pass on jax 0.10.1, float64. |
The non-Hermitian eigenvectors are returned with stop_gradient, which also zeroes the gradient of an otherwise gauge-invariant spectral projector X @ Y^H. Correct the docstring (which previously advertised projectors as differentiable) to state only eigenvalues are differentiable, and lock the documented limitation with a projector-gradient regression test. Co-authored-by: Xing Zhang <fishjojo@users.noreply.github.com>
1 participant
Summary
Adds
pyscfad/lib/davidson_lite.pywith fullyjax-jittable Davidson eigensolvers that are differentiable to arbitrary order (forward and reverse mode), implementing the request in #116.Instead of differentiating through the Davidson iterations, the converged solution is differentiated via the implicit function theorem using
jax.lax.custom_root, mirroring theSCFLitepattern inpyscfad/scf/hf_lite.py.eigh— Hermitian / real-symmetric (and complex-Hermitian) operators. Fully implicitly differentiable. The (near-)degenerate response is projected out viadeg_thresh, following the masking strategy inpyscfad/backend/_jax/lax/linalg.py::_eigh_gen_jvp_rule.eig— general non-Hermitian operators. The adjoint uses an explicit transposed operatoraopT(Aᴴ) to recover the biorthonormal left eigenvectors and build the oblique projectors.Key properties (per the issue discussion):
X = -g((0, 1))), which is what makes higher-order derivatives correct.nrootsso the solver stays jittable (callers request enough roots to cover any boundary degeneracy).lax.while_loop, static shapes, batchedaop(V), thick restart viadynamic_update_slice.aop(so a complex operator applied to a real guess works).The existing
davidson1and itstdscf/fcicall sites are left untouched — this lands as a new module with no migration.Validation
tests/test_davidson_lite.py(10 tests, all passing locally on jax 0.10.1, float64):eigh: forward vsjnp.linalg.eigvalsh; residual + orthonormality;jit; Hellmann–Feynman gradient;check_grads(order=2, modes=['fwd','rev']); eigenvector-gradient vs denseeigh; degenerate spectrum; complex-Hermitian forward + gradient.eig: forward + biorthonormal left vectors (Yᴴ X = I);jit;check_grads(order=2, fwd+rev)for the eigenvalue.Note on non-Hermitian targeting
Lowest-
Nselection foreigrelies on the operator being diagonally dominant (the CI/TDDFT regime), so the diagonal preconditioner and unit-vector guess steer toward the lowest real-part roots. For strongly non-normal operators the converged pairs are still exact eigenpairs but may not be the lowest ones; supply a goodx0/adiagin that case. This is documented in theeigdocstring.Closes #116.
Test plan
tests/test_davidson_lite.pyGenerated with Claude Code