🔪 The Sharp Bits 🔪#

Compatible devices#

  • NVIDIA GPU: Heavily tested, best performance

  • AMD GPU: To be supported in the future

  • CPU: Supported but not heavily tested

  • Google TPU: Not supported

Global system#

Quantax stores the system geometry and Hilbert space — the Sites (or Lattice) — as a process-global constant, retrieved anywhere via quantax.get_sites() and quantax.get_lattice(). This differs from most NQS packages, where the system is an ordinary object passed around explicitly.

The consequence is that only one Sites should exist per Python process. Constructing a second one does not raise an error; it merely warns and silently overwrites the first, which can lead to hard-to-debug behavior. This is especially easy to trigger in notebooks by re-running a cell that builds the system.

To switch to a different system, restart the Python interpreter (or the notebook kernel).

Randomness#

Quantax keeps a single global PRNG key, replicated across all devices. Set it with quantax.set_random_seed() (the default seed is 42) for reproducible runs, and draw fresh subkeys with quantax.get_subkeys().

Because get_subkeys reads and updates this global key, it is not jittable. Call it outside jitted functions and pass the resulting keys in as arguments.

Data precision#

By default, JAX

  • disables double precision, so float64/complex128 silently fall back to float32/complex64;

  • uses TF32, a reduced-precision format, for the matmuls and convolutions of float32 inputs on supported GPUs, whenever possible.

Quantax does NOT change these defaults. Most computations work well under this setting, but be careful in the following cases.

  • Matrix inversion / linear solves. Several solvers in quantax/optimizer/solver.py switch to FP64 (with jax.enable_x64()) or FP32 matmuls (precision='highest') locally to keep the inverse accurate.

  • Local updates. Higher precision might be required in local updates to avoid error accumulation.

To enable double precision globally, call jax.config.update("jax_enable_x64", True). Alternatively, call quantax.set_default_dtype(jnp.float64) to set jnp.float64 as the default data type in Quantax, which also turns on jax_enable_x64. To disable TF32, call jax.config.update("jax_default_matmul_precision", "float32").

The variational models in quantax.model deliberately ignore the Quantax default dtype and pick their own (usually float32) for efficiency, so quantax.set_default_dtype() does not change them.

To JIT or not to JIT#

jax.jit greatly accelerates JAX functions, but it imposes constraints on the jitted function and can make the code harder to read.

In Quantax, our strategy is to apply jax.jit only to functions that are likely to be the bottleneck of a simulation, and to keep the rest unjitted for flexibility. For instance, quantax.operator.Operator.expectation() is left unjitted so that it can compute the expectation of arbitrary states, while the forward pass of quantax.state.Variational — the actual bottleneck when computing expectation values — is jitted.

Wavefunction overflow#

Many-body wavefunctions often span a wide range of magnitudes, even beyond the limit of FP64. In many packages, e.g. NetKet and jVMC, the wavefunction amplitudes are expressed by \(\log \psi\) instead of \(\psi\) to avoid overflow (and underflow).

In Quantax, we don’t utilize this approach for the following reasons.

  1. Ground-state wavefunctions are often real-valued with signs. Expressing them using \(\log \psi\) forces one to perform optimization in the complex space.

  2. \(\log\psi\) is discontinuous at \(\psi=0\), which might cause stability issues.

Instead, Quantax keeps \(\psi\) but factors out its magnitude into a separate, overflow-safe scale. Two custom array types in quantax.utils implement this, each a PyTree with two leaves:

  • LogArray — value = sign * exp(logabs), with sign being \(\pm 1\) or a complex phase and logabs the real log-magnitude.

  • ScaleArray — value = significand * exp(exponent), with exponent a real normalization factor pulled out to keep significand in range.

The sign/phase stays attached to the amplitude, so a real sign-structured wavefunction remains real and continuous at \(\psi=0\). Arithmetic is overloaded to act in the scaled space (e.g. multiplying two LogArrays adds the logabs, never forming the overflowing product); convert back to jax.Array with arr.value() or jnp.asarray(arr).

For details of constructing neural wavefunctions with LogArray and ScaleArray, see tutorial build your network.

Sharding#

Quantax uses JAX’s auto sharding by default. For details, see the JAX guide Distributed arrays and automatic parallelization.

The mesh has two axes, ("process", "device"), with shape (jax.process_count(), jax.local_device_count()).

Some arrays, such as neural-network weights, are replicated across all devices (see get_replicated_sharding()). Others, such as Monte Carlo samples, are distributed over the devices (see get_distributed_sharding()).

HPC usage#

There are 3 ways to submit JAX jobs on HPC. Here we provide the caveats in the python script and the slurm scripts on the NERSC cluster as examples. In all cases, the users only need a few lines of code to control the distributed initialization of JAX, while the Quantax API works the same way as on a single device.

Single-node, only one process#

In the single-node case, one doesn’t have to do anything special in the python script. The slurm script is given below.

#!/bin/bash -l
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-task=<number_of_gpus>
#SBATCH --output=gpu-out.%j
#SBATCH --error=gpu-err.%j
#SBATCH --time=<time_of_simulations>
#SBATCH --constraint=gpu
#SBATCH --qos=regular
#SBATCH --account=<account_name>

source <path_to_quantax_env>/bin/activate

srun --cpu-bind=none python -u <path_to_python_script>

Multi-node, one process per GPU#

#!/bin/bash -l
#SBATCH --nodes=<number_of_nodes>
#SBATCH --ntasks-per-node=4  # 4 GPUs / node on NERSC
#SBATCH --gpus-per-task=1
#SBATCH --gpu-bind=none
#SBATCH --mem=0
#SBATCH --output=gpu-out.%j
#SBATCH --error=gpu-err.%j
#SBATCH --time=<time_of_simulation>
#SBATCH --constraint=gpu
#SBATCH --qos=regular
#SBATCH -c 32
#SBATCH --account=<account_name>

source <path_to_quantax_env>/bin/activate

# Specific to multi-node jobs on NERSC
module load nccl/2.24.3

# Some JAX versions fail without turning off xla_gpu_shard_autotuning
export XLA_FLAGS="--xla_gpu_shard_autotuning=false"

srun --cpu-bind=cores python -u <path_to_python_script>

One should configure the distributed initialization of JAX in the beginning of the python script.

import jax
jax.distributed.initialize()

...

Multi-node, one process per node#

#!/bin/bash -l
#SBATCH --nodes=<number_of_nodes>
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=4  # 4 GPUs / node on NERSC
#SBATCH --gpu-bind=none
#SBATCH --mem=0
#SBATCH --output=gpu-out.%j
#SBATCH --error=gpu-err.%j
#SBATCH --time=<time_of_simulation>
#SBATCH --constraint=gpu
#SBATCH --qos=regular
#SBATCH -c 128
#SBATCH --account=<account_name>

source <path_to_quantax_env>/bin/activate

# Specific to multi-node jobs on NERSC
module load nccl/2.24.3

# Some JAX versions fail without turning off xla_gpu_shard_autotuning
export XLA_FLAGS="--xla_gpu_shard_autotuning=false"

srun --cpu-bind=none python -u <path_to_python_script>

jax.distributed.initialize often fails to handle this case automatically. The following code resolves it.

import os
import jax

gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])

jax.distributed.initialize(
    num_processes=int(os.environ["SLURM_NTASKS"]),
    process_id=int(os.environ["SLURM_PROCID"]),
    local_device_ids=list(range(gpus_per_node)),
)

...