quantax.state.Variational#

class quantax.state.Variational#

Bases: State

Variational state. This is a wrapper of a jittable variational ansatz. The variational model should be given as an equinox.Module. For details of Equinox, see this documentation.

Warning

There are many intermediate values stored in the class, so most functions like __call__ in this class are non-jittable.

Warning

Many quantities are only computed once in the initialization. Please don’t change the private attributes unless an update function is provided. One can define a new state if some changes are necessary.

Parameters:

model – Variational model. Should be an equinox.Module.
param_file – File for loading parameters which is saved by save or equinox.tree_serialise_leaves, default to not loading parameters.
symm – Symmetry of the network, default to Identity. Denoting the network output as \(f(s)\) and symmetry elements as \(T_i\) with characters \(\omega_i\), the wave function is given by \(\psi(s) = \sum_i \omega_i \, f(T_i s) / n_{symm}\)
max_parallel –
The maximum chunk size allowed per device. Specifying a limited value is important for avoiding memory overflow. For many hamiltonians, this also helps to improve the efficiency by keeping a constant amount of forward pass and avoiding re-jitting. The allowed input formats are:
- None:
  No chunk size (default).
- int:
  The same chunk size for all forward and backward passes.
- Tuple[int, int]:
  (forward chunk, backward chunk)
- Tuple[int, int, int]:
  (forward chunk, backward chunk, ref chunk) Ref chunk is the chunk size used in init_internal and ref_forward. If ref chunk is not specified by this format, it defaults to the forward chunk size.
use_ref – Whether ref_forward will be used when the model is a RefModel. When the model is not a RefModel, this argument has no effect. Default to True.

__getitem__(basis_ints: NDArray | Array) → NDArray | Array | LogArray | ScaleArray#

Evaluate the wave function \(\psi(s) = \left<s|\psi\right>\) by state[s]

Parameters:: basis_ints – A batch of basis integers

__matmul__(other: State) → complex#: Compute the contraction \(\left< \psi|\phi \right>\) by self @ other. This is implemented by converting self and other to DenseState.

__call__(s: NDArray | Array) → NDArray | Array | LogArray | ScaleArray#

Evaluate the wavefunction \(\psi(s) = \left<s|\psi\right>\).

Parameters:: s – Spin/fermion configurations s with entries \(\pm 1\).

Warning

This function is not jittable.

Note

The returned value is \(\psi(s)\) instead of \(\log\psi(s)\).

property symm: Symmetry#: Symmetry of the state

property use_ref: bool#: Whether to use reference implementation for updates

property model: Callable[[Array], NDArray | Array | LogArray | ScaleArray]#: The variational model used in the variational state.

property holomorphic: bool#: Whether the variational state is holomorphic.

property forward_chunk: int | None#: The maximum chunk size of forward pass allowed per device.

property backward_chunk: int | None#: The maximum chunk size of backward pass allowed per device.

property ref_chunk: int | None#: The maximum chunk size of ref_forward allowed per device.

property nparams: int#: Number of total parameters in the variational state.

property dtype: str | type[Any] | dtype | SupportsDType#: The parameter data type of the variational state.

property vs_type: VS_TYPE#: The type of variational state.

fast_forward(s: Array) → NDArray | Array | LogArray | ScaleArray#

Evaluate the wavefunction \(\psi(s) = \left<s|\psi\right>\). This function assumes s to be in good shape and sharding for speedup.

Parameters:: s – Spin/fermion configurations s with entries \(\pm 1\)

init_internal(s: Array) → tuple[NDArray | Array | LogArray | ScaleArray, PyTree]#: Return the wavefunction and initial internal values for the given input s.

property required_update_modes: tuple[str, ...]#: The required update modes for accelerated ref_forward pass.

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: Literal[False] = False) → NDArray | Array | LogArray | ScaleArray#

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: Literal[True]) → tuple[NDArray | Array | LogArray | ScaleArray, PyTree]

Compute the forward pass given reference internal state of the model.

Parameters:

s – Input states s with entries \(\pm 1\).
s_old – The old states before the updates, with entries \(\pm 1\).
update_mode – The update modes required by the model.
internal – The internal state of the model, which is initialized by init_internal.
return_update – Whether to return the updated internal state along with the wave function.

Returns:

The output wave function \(\psi(s)\), or a tuple of the wave function and the updated internal state of the model if return_update is True.

segment_ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], idx_segment: Array, internal: PyTree) → NDArray | Array | LogArray | ScaleArray#

Compute the forward pass with segments given reference internal state of the model. This method is usually used in the computation of local energy.

Parameters:

s – Input states s with entries \(\pm 1\).
s_old – The old states before the updates, with entries \(\pm 1\).
update_mode – The update modes required by the model.
idx_segment – The indices of the segment to be updated, which is used to slice s_old and internal.
internal – The internal state of the model, which is initialized by init_internal.

Returns:

The output wave function \(\psi(s)\).

jacobian(s: Array) → Array#

Compute the jacobian matrix \(\frac{1}{\psi} \frac{\partial \psi}{\partial \theta}\). See VS_TYPE for the definition of jacobian for different kinds of networks.

Parameters:: s – The input spin/fermion configurations with entries \(\pm 1\).
Returns:: A 2D jacobian matrix with the first dimension for different inputs and the second dimension for different parameters. The order of parameters are the same as get_params_flatten.

precompile(mesh: Mesh | None = None, forward_batches: Sequence[int] | int = (), backward_batches: Sequence[int] | int = ()) → None#

Ahead-of-time compile the expensive jitted functions of this state and write the executables into the persistent compilation cache.

Two usages:

precompile(...) without a mesh compiles for the devices of the current run, so the compilation cost is paid here instead of at the first forward/backward call later in the program (which then loads the executables from the persistent cache).
With a compile-only mesh from make_precompile_mesh, this compiles for a different (typically larger) device topology. A real run on that topology with the same environment loads the executables from the shared cache instead of compiling, so a cheap single-node job can absorb the compilation cost of a multi-node job.

Parameters:

mesh – The ("process", "device") mesh describing the target topology, default to make_mesh spanning the devices of the current run.
forward_batches – Global batch sizes to compile the jitted chunked forward (fast_forward) for; a single int is treated as one batch size. Jitted functions are compiled once per batch shape, and the batches of the target run are determined by its samplers: the chain numbers of the (sub-)samplers, and the chunk-multiple update batches of chunked sweeps. They can be read off a cheap pilot run of the target script with jax.config.update("jax_log_compiles", True) from the jit(chunked_f) log lines; the global batches are topology-independent, so a 1-GPU pilot with the target’s global number of samples reports the correct values.
backward_batches – Global batch sizes to compile the jacobian (jacobian) for, typically the global number of samples; a single int is treated as one batch size.

The forward used by __call__() and observable estimation is always compiled (its batch is fixed to forward_chunk * mesh.size by the chunking) when max_parallel is set. max_parallel must equal the value used by the target run.

Warning

The persistent-cache key includes the jaxlib version, XLA flags, GPU model, and CUDA version. Run the precompiling job on one node of the same cluster with the same environment and XLA_FLAGS as the target run, with jax_compilation_cache_dir pointing to a filesystem shared with it.

Note

Compiling for a compile-only mesh first executes a small real forward/backward warm-up on the current devices, with the target run’s per-device batches, to record cuDNN/GEMM autotuning results in the XLA per-fusion autotune cache; the ahead-of-time compilation reuses them instead of profiling (which is impossible without real devices). This requires the default jax_persistent_cache_enable_xla_caches setting, which keeps the autotune cache enabled alongside the persistent cache.

Note

Small helper functions (samplers, chunking utilities, solvers) still compile natively in the target run within seconds, and Python tracing is still paid by every process of the target run.

Note

The jitted functions of this state stay bound to the given mesh afterwards. With a compile-only mesh, executing them raises an error, so the state cannot be used for computation after precompiling for a different topology.

Note

Reference forwards (use_ref) are not covered yet.

partition(model: PyTree = None) → tuple[PyTree, PyTree]#

Split the variational model into two pytrees, one containing all parameters and the other containing all other elements, similar to partition in Equinox.

Parameters:: model – The model to be split, default to be the variational model in the variational state.

combine(params: PyTree, others: PyTree) → PyTree#

Combine two pytrees, one containing all parameters and the other containing all other elements, into one variational model. This is similar to combine in Equinox.

Parameters:

params – The pytree containing only parameters.
others – The pytree containing other elements.

get_params_flatten() → Array#: Obtain a flattened 1D array of all parameters.

get_params_unflatten(params: Array) → PyTree#: Obtain the parameters pytree from a flattened 1D array of all parameters.

update(step: Array) → None#

Update the variational parameters of the state as \(\theta' = \theta - \delta\theta\).

Parameters:: step – The update step \(\delta\theta\).

Note

The update direction is \(-\delta\theta\) instead of \(\delta\theta\).

save(file: str | Path | BinaryIO) → None#: Save the variational model in the given file. This file can be used to load when initializing Variational.

to_netket_model() → Module#

Convert the state to an equinox.Module compatible with NetKet. NetKet natively accepts equinox modules as the variational ansatz of an nk.vqs.MCState, so the returned module can be passed directly to NetKet to measure observables.

The module takes spin configurations with entries \(\pm 1\) (NetKet’s convention for netket.hilbert.Spin) and returns \(\log\psi\). Whatever the underlying quantax model outputs (jax.Array, LogArray, or ScaleArray), it is converted to \(\log\psi\). The output is always complex so that sign-structured or complex wavefunctions are represented correctly (see the wavefunction overflow section of the sharp_bits tutorial).

Warning

Training the generated state in NetKet is probably unstable, but the state can be reliably used to measure observables.