A variational Monte Carlo engine built directly on JAX. Hamiltonians are composed from spin and fermionic operator terms, any Flax module can serve as the ansatz, and the same program runs on a single CPU and across a multi-GPU cluster.
Spin-½ lattice models, benchmarked up to ~0 spins.
Fermionic lattice models, benchmarked up to 0 electrons.
A single network trained across a family of Hamiltonians rather than one per coupling.
A ground-state calculation consists of three function calls in an explicit Python loop: configurations are sampled, the energy is estimated on that sample, and the parameters are updated. Each call returns new values rather than modifying its arguments; there is no driver object and no callback mechanism.
A move rule defines a Markov chain over configurations. sample advances every chain and returns the resulting configurations together with their log-amplitudes.
state, log_amps, acceptance = sample(nsweeps, state, action, keys, wf)
Chains are sharded across the available JAX devices, so the call is the same on one CPU and on many GPUs.
A Hamiltonian is a Python callable assembled from operator terms. compute_expectation applies it to the sampled configurations and returns the local energies together with ⟨E⟩ and ⟨E²⟩.
E_L, e_mean, e2_mean = compute_expectation(H, wf, state, log_amps)
No matrix representation is constructed, so memory scales with the sample rather than with the dimension of the Hilbert space.
An optimizer maps the local energies to a parameter update. Applying the update returns a new wavefunction rather than modifying the existing one, so a run is a sequence of values that can be checkpointed or compared at any step.
updates, opt_state = optimizer(E_L, opt_state, state, wf) wf = wf.apply_gradients(updates, lr)
Learning-rate schedules, stopping criteria and additional observables are expressed directly in the loop rather than registered through hooks.
Each of these is a pure function over JAX pytrees. Any Flax module can be used as the ansatz, and the loop can be passed to jit or vmap without further wrapping.
tachys is installed with pip install tachys and requires JAX.