dftax writes Kohn-Sham DFT (integrals, quadrature grid, exchange-correlation, the SCF loop) as a single differentiable JAX program. Once the energy is one autodiff-traceable function, the derivatives quantum chemistry cares about are just calls to jax.grad. Forces are −∂E/∂R; the Fock matrix is sym(∂E/∂P), so no exchange-correlation potential is coded by hand; response properties are higher derivatives of the same function.

forces(mol, xc, res)         # −∂E/∂R    nuclear forces
hessian(mol, xc)             # ∂²E/∂R²   vibrational frequencies
polarizability(mol, xc)      # −∂²E/∂F²
alchemical_deriv(mol, xc)    # ∂E/∂Z

Converged quantities are differentiated through the SCF fixed point (CPHF), so a gradient of a converged calculation costs one linear solve rather than an unrolled solver. And because it is ordinary JAX, the whole engine jits, vmaps over geometries, and runs on GPU; a DFT calculation drops into a training loop or an end-to-end functional fit like any other differentiable block.

The runtime carries no libcint, libxc, or Maple: basis sets come from the Basis Set Exchange at setup, the molecular grid is built natively, and PySCF appears only as a test-time reference.

Installation

pip install dftax            # CPU
pip install dftax[cuda12]    # + CUDA-12 jaxlib (Linux GPU)

Quick start

import jax
jax.config.update("jax_enable_x64", True)   # DFT energies want float64

from dftax import KS, Molecule, scf
from dftax.energy.xc import PBE   # or functional("..."): LDA, PBE0, B3LYP,
                                  # CAM-B3LYP, wB97X(-V), wB97M-V, r2SCAN

mol = Molecule.from_xyz("O 0 0 0; H 0.757 0.587 0; H -0.757 0.587 0", "sto-3g")
res = scf(KS(mol, PBE()))            # DIIS SCF -> KSResult
print(res.e_tot, res.converged)     # total energy (Ha), convergence flag

KS(system, xc, *, grid=, coulomb=, spin=, mesh=, dispersion=) assembles the energy functional; the verbs scf, newton, roks, and minimize solve it, and forces, dipole, polarizability, and the rest read properties from it. Every choice is a value passed to the builder, not a global flag: a nonzero spin (given as 2S) runs the spin-polarized α/β path, dispersion=d3bj() or d4() adds a dispersion correction, scf(ks, accel=adiis()) accelerates hard SCF cases, and scf(ks, smearing=fermi()) gives fractional occupations and the Mermin free energy for metallic or degenerate systems.

Scaling up

The exact 4-center ERI is O(N⁴): the small-system reference the rest is checked against. From there each backend is one more value on the builder. Density fitting (df(...)) takes the Coulomb cost to O(N³); streaming holds the tensors at O(N²) memory; Schwarz screening trims the per-iteration work; and mesh=mesh() shards the grid and the density-fitting tensor across GPUs. Every collective differentiates, so SCF, minimization, and forces run unchanged.

from dftax import KS, df, mesh, scf
from dftax.energy.xc import PBE0

scf(KS(mol, PBE0(), coulomb=df("def2-universal-jkfit"), mesh=mesh()))

The same mesh spans several nodes: launch the script once per node and call distributed() first, and the auxiliary slabs shard across every GPU of the allocation (scripts/gpu/distributed.sbatch is a working template).

Accuracy

Against PySCF on water / sto-3g (see scripts/bench/BENCHMARKS.md):

LDA, PBE, B3LYP and PBE0 reproduce libxc to machine precision (an earlier ~1e-5 Ha PBE gap traced to one truncated constant, now exact). The later rungs were ported against their references directly: r2SCAN matches libxc to ~1e-9, ωB97X-V (including its VV10 term) matches a full PySCF solve to ~2e-13 on matched grids, ωB97M-V matches libxc pointwise to machine precision and a full PySCF solve to ~1e-10, and the D3(BJ)/D4 dispersion models match tad-dftd3/tad-dftd4 to machine precision. Analytic forces match finite differences to about 4e-8 Ha/Bohr, the CPHF polarizability matches finite-field to about 1e-4, and vibrational frequencies match PySCF to a few cm⁻¹.

Limitations

• Orbital bases run up to i shells (l=6), i.e. through cc-pV5Z/6Z; use the density-fitting backend at those sizes (the exact 4-center path is correct but impractically slow there, as in conventional codes).
• The materialized exact-ERI path hits a GPU memory ceiling near nao ≈ 70; use density fitting (the default, with a device-aware memory policy), streaming, and screening at scale (exact-exchange compute stays O(N⁴) regardless).
• VV10 (ωB97X-V) needs the materialized XC grid: the pair quadrature does not stream or shard yet.
• GPU correctness is validated interactively on an A100 rather than in CPU CI, and large-N GPU throughput is not yet benchmarked.

Documentation

Docs, tutorials, and executed example notebooks live at https://andresguzco.github.io/dftax/; the same examples are runnable scripts under examples/.

See also

dftax builds on Equinox and Optax, and shares ground with other differentiable electronic-structure codes: MESS and D4FT (in JAX), DQC (PyTorch), and PySCFad. It leans on PySCF as its test-time reference and aims to be a good neighbor to it, not a replacement.

Citation

@software{dftax,
  author  = {Guzm{\'a}n-Cordero, Andr{\'e}s},
  title   = {dftax: a differentiable Kohn-Sham DFT engine in JAX},
  url     = {https://github.com/andresguzco/dftax},
  version = {0.8.1},
  year    = {2026},
}

License

Apache-2.0.