Welcome to the Quantax documentation!

Quantax is a Python package for flexible and scalable neural quantum states built on JAX

Installation

Requires Python 3.10+, JAX 0.6.1+

First, ensure that a correct JAX version is installed. For details, check JAX Installation.

For a direct installation of full functionality (recommended in most cases),

pip install quantax[full]

For a minimal installation,

pip install quantax

Tutorials

The tutorials below guide you through the main features of Quantax. You don’t have to follow them one by one. Feel free to jump to the ones that interest you most.

• Quick Start
• Exact diagonalization
• Build your network
• Samples and Measurement
• Square J1-J2 model
• Triangular Heisenberg model
• Fermion mean field
• Neural Jastrow
• Real-time dynamics
• Local updates

Examples

The examples below reproduce the main results of several important NQS papers.

• Restricted Boltzmann machine
• Time-dependent variational principle
• Minimum-norm stochastic reconfiguration
• Neural network backflow

API

global_defs	Modifiers and utility functions of global constants
sites	Geometry and particles of the quantum system
symmetry	Symmetry sector of the studied system
operator	Quantum operators on the Hilbert space
nn	Network components
model	Variational wavefunctions
state	Quantum states
sampler	Samplers for generating configurations
optimizer	Optimizers for training wavefunctions
utils	Utility functions

Papers

• Related packages
• Research papers