netket.optimizer

This module provides some optimisers, implementations of the Quantum Geometric Tensor and preconditioners such as SR.

Optimizers

Optimizers in NetKet are simple wrappers of optax optimizers. If you want to write a custom optimizer or use more advanced ones, we suggest you have a look at optax documentation.

Check it out for up-to-date information on available optimisers.

Warning

Even if optimisers in netket.optimizer are optax optimisers, they have slightly different names (they are capitalised) and the argument names have been rearranged and renamed. This was chosen in order not to break our API from previous versions

In general, we advise you to directly use optax, as it is much more powerful, provides more optimisers, and it’s extremely easy to use step-dependent schedulers.

Adam	Adam Optimizer.
AdaGrad	AdaGrad Optimizer.
Sgd	Stochastic Gradient Descent Optimizer.
Momentum	Momentum-based Optimizer.
RmsProp	RMSProp optimizer.

Preconditioners

This module also provides an implementation of the Stochastic Reconfiguration/Natural gradient preconditioner.

SR	Stochastic Reconfiguration or Natural Gradient preconditioner for the gradient.

Quantum Geometric Tensor

It also provides the following implementation of the quantum geometric tensor:

qgt.QGTAuto	Automatically select the 'best' Quantum Geometric Tensor computing format according to some rather untested heuristic.
qgt.QGTOnTheFly	Lazy representation of an S Matrix computed by performing 2 jvp and 1 vjp products, using the variational state's model, the samples that have already been computed, and the vector.
qgt.QGTJacobianPyTree	Semi-lazy representation of an S Matrix where the Jacobian O_k is precomputed and stored as a PyTree.
qgt.QGTJacobianDense	Semi-lazy representation of an S Matrix where the Jacobian O_k is precomputed and stored as a dense matrix.

Dense solvers

And the following dense solvers for Stochastic Reconfiguration:

solver.cholesky_with_fallback	Solve the linear system using a Cholesky factorisation, automatically falling back to pinv_smooth() when the result contains NaN or Inf.
solver.cholesky	Solve the linear system using a Cholesky Factorisation.
solver.LU	Solve the linear system using a LU Factorisation.
solver.pinv	Solve the linear system using jax's implementation of the pseudo-inverse.
solver.pinv_smooth	Solve the linear system by building a pseudo-inverse from the eigendecomposition obtained from jax.numpy.linalg.eigh().
solver.solve	Solve the linear system.
solver.svd	Solve the linear system using Singular Value Decomposition.

Solver combinators

solver.nan_fallback

Creates a solver that transparently falls back to a more robust solver when the primary produces NaN or Inf.

Multi-GPU solvers

And the following dense solvers parallelize among multiple GPUs by making use of JAXMg. They are not really much faster than the ones above, but allow you to use ~100k samples because they keep the NTK sharded across the multiple GPUs.

Warning

Those solvers require jaxmg >= 1.4.1, which is not installed with NetKet. Install it with the CUDA version matching your jax (pip install 'jaxmg[cuda12]' or pip install 'jaxmg[cuda13]'), which also installs NVIDIA’s cuSOLVERMp library. Every jaxmg release pins a narrow range of jax versions, so installing it may change the version of jax you use.

As JAXMg builds upon NVIDIA’s cuSOLVERMp, it requires one python process per GPU: run your script with djaxrun (or call jax.distributed.initialize() yourself) such that every process sees a single GPU. Both single-node and multi-node setups are supported.

solver.cholesky_distributed	Solve the linear system using a distributed Cholesky factorization.
solver.pinv_smooth_distributed	Solve the linear system using a distributed pseudo-inverse from eigendecomposition.