eigh 0.2.0

Differentiable eigenvalue decomposition with JAX (CPU/GPU)

Differentiable Generalized Eigenvalue Decomposition

Standalone implementation of differentiable eigenvalue decomposition with CPU (LAPACK) and GPU (cuSOLVER) backends. Extracted from pyscfad.

Wheels on PyPI: https://pypi.org/project/eigh/ — Linux (manylinux_2_28, x86_64) and macOS (x86_64, arm64), Python 3.10–3.12. GPU path (cuSOLVER) is tested locally; CI runs CPU tests only.

Windows: no prebuilt wheel. The pure-JAX solvers in src/jax/ (e.g. safe_generalized_eigh, subspace_generalized_eigh, stable_generalized_eigh) work out-of-the-box — pip install jax numpy scipy and import directly from that module. The fast LAPACK/cuSOLVER-backed eigh / eigh_gen kernels (and therefore stable_eigh_pyscfad / stable_eigh_gen_pyscfad) require building from source against a local BLAS/LAPACK.

Features

• Generalized Problems: A @ V = B @ V @ diag(W), etc.
• JAX Integrated: Full support for jit, vmap, grad, and jvp.
• High Performance: Optimized LAPACK (CPU) and cuSOLVER (GPU) kernels.
• Precision: float32/64 and complex64/128.
• Degeneracy Handling: Configurable deg_thresh for stable gradients.

Installation & Quick Start

# Install from source
pip install .

# For GPU support in this environment
pip install .[cuda-local]

Usage Example

import jax
import jax.numpy as jnp
# Gen. eigensolver from PySCFAD
from eigh import eigh

jax.config.update("jax_enable_x64", True)
A = jnp.array([[2., 1.], [1., 2.]])
w, v = eigh(A) # Standard
grad = jax.grad(lambda A: eigh(A)[0].sum())(A) # Differentiable

Benchmarks

Forward/backward scaling vs. matrix size, and gradient stability as eigenvalues approach degeneracy — for the JAX eigensolvers in src/jax/. See benchmarks/suite/ for the scripts.

API Reference

• eigh(a, b=None, *, lower=True, eigvals_only=False, type=1, deg_thresh=1e-9) Scipy-compatible interface. type supports 1: A@v=B@v@λ, 2: A@B@v=v@λ, 3: B@A@v=v@λ.
• eigh_gen(a, b, *, lower=True, itype=1, deg_thresh=1e-9) Lower-level generalized solver.

Degenerate Eigenvalues & Gradients

Individual eigenvalue gradients are ill-defined for degenerate (repeated) eigenvalues. However, symmetric functions (like sum, var, trace) have stable gradients. The deg_thresh parameter (default 1e-9) masks divisions by near-zero gaps to maintain stability.

JAX Eigensolvers (src/jax/)

A collection of differentiable generalized eigensolvers with different strategies for handling degenerate eigenvalues in reverse-mode gradients. Useful for training pipelines where degeneracies are common.

Recommended: stable_eigh_pyscfad / stable_eigh_gen_pyscfad from generalized_eigensolver_pyscfad.py. These are the solvers promoted by this package — they wrap the fast LAPACK/cuSOLVER-backed eigh / eigh_gen kernels with a Lorentzian-broadened custom VJP for stable gradients at (near-)degeneracies. The other solvers below are provided for comparison and experimentation.

Solver	File	Strategy	Gradient-safe at degeneracies
stable_eigh_pyscfad / stable_eigh_gen_pyscfad ⭐	generalized_eigensolver_pyscfad.py	Wraps eigh/eigh_gen (LAPACK/cuSOLVER) with Lorentzian-broadened VJP [2]	Yes — recommended
standard_eig	generalized_eigensolver.py	scipy.linalg.eigh (CPU, non-diff)	N/A
jax_eig	generalized_eigensolver.py	Plain Cholesky + jnp.linalg.eigh	No
generalized_eigh	generalized_eigensolver.py	Symmetrized Cholesky with SPD shift	No (standard VJP)
degen_eigh	generalized_eigensolver.py	Custom VJP: mask degenerate F_ij via thresholding [1,3]	Yes, for symmetric-subspace losses
safe_generalized_eigh	generalized_eigensolver.py	Cholesky + degen_eigh	Yes
subspace_eigh	generalized_eigensolver.py	Custom VJP: Lorentzian broadening F/(F²+ε²) [2]	Yes
subspace_generalized_eigh	generalized_eigensolver.py	Symmetry-breaking perturbation + subspace_eigh [2,4]	Yes
stable_eigh / stable_generalized_eigh	generalized_eigensolver_stable.py	Pure-JAX Cholesky + Lorentzian-broadened VJP [2]	Yes

References (verified)

• [1] Kasim, M. F., & Vinko, S. M. Learning the exchange–correlation functional from nature with fully differentiable density functional theory. Phys. Rev. Lett. 127, 126403 (2021). https://doi.org/10.1103/PhysRevLett.127.126403
• [2] Colburn, S., & Majumdar, A. Inverse design and flexible parameterization of meta-optics using algorithmic differentiation. Communications Physics 4, 54 (2021). https://doi.org/10.1038/s42005-021-00568-6
• [3] JAX Issue #2748 — Differentiable eigh with degeneracies. https://github.com/jax-ml/jax/issues/2748
• [4] JAX Issue #5461 — Stable generalized eigh. https://github.com/jax-ml/jax/issues/5461

Development & Testing

• Requirements: CMake 3.18+, C++17, JAX, NumPy, LAPACK/CUDA.
• Tests:

  pytest tests/test_eigh.py     # Core functionality
  pytest tests/test_eigh_gen.py # Generalized itypes
  pytest tests/test_eigh_jit.py # JIT & vmap
  

• GPU Setup:

  source setup_gpu_env_clean.sh
  ./run_gpu.sh python example_simple.py
  

License & Citation

Apache License 2.0. If used in research, please cite:

@software{pyscfad,
  author = {Zhang, Xing},
  title = {PySCFad: Automatic Differentiation for PySCF},
  url = {https://github.com/fishjojo/pyscfad},
  year = {2021-2025}
}

@software{sokolov2026eigh,
  title={Eigh: Differentiable eigenvalue decomposition with jax (cpu/gpu)},
  author={Sokolov, Igor},
  url={https://github.com/Brogis1/eigh},
  year={2026}
}

@article{sokolov2026xc,
  title = {Quantum-enhanced neural exchange-correlation functionals},
  author = {Sokolov, Igor O. and Both, Gert-Jan and Bochevarov, Art D. and Dub, Pavel A. and Levine, Daniel S. and Brown, Christopher T. and Acheche, Shaheen and Barkoutsos, Panagiotis Kl. and Elfving, Vincent E.},
  journal = {Phys. Rev. A},
  volume = {113},
  issue = {1},
  pages = {012427},
  numpages = {24},
  year = {2026},
  month = {Jan},
  publisher = {American Physical Society},
  doi = {10.1103/m51l-fys2},
  url = {https://link.aps.org/doi/10.1103/m51l-fys2}
}