netket.optimizer.qgt.QGTJacobianPyTree

netket.optimizer.qgt.QGTJacobianPyTree#

netket.optimizer.qgt.QGTJacobianPyTree(vstate=None, *, mode=None, holomorphic=None, diag_shift=None, diag_scale=None, rescale_shift=None, chunk_size=None, **kwargs)[source]#

Semi-lazy representation of an S Matrix where the Jacobian O_k is precomputed and stored as a PyTree.

The matrix of gradients O is computed on initialisation, but not S, which can be computed by calling to_dense. The details on how the ⟨S⟩⁻¹⟨F⟩ system is solved are contained in the field sr.

Numerical estimates of the QGT are usually ill-conditioned and require regularisation. The standard approach is to add a positive constant to the diagonal; alternatively, Becca and Sorella (2017) propose scaling this offset with the diagonal entry itself. NetKet allows using both in tandem:

\[S_{ii} \mapsto S_{ii} + \epsilon_1 S_{ii} + \epsilon_2;\]

\(\epsilon_{1,2}\) are specified using diag_scale and diag_shift, respectively.

Parameters:
  • vstate – The variational state

  • mode (Optional[str]) – “real”, “complex” or “holomorphic”: specifies the implementation used to compute the jacobian. “real” discards the imaginary part of the output of the model. “complex” splits the real and imaginary part of the parameters and output. It works also for non holomorphic models. holomorphic works for any function assuming it’s holomorphic or real valued.

  • holomorphic (Optional[bool]) – a flag to indicate that the function is holomorphic.

  • diag_scale – Fractional shift \(\epsilon_1\) added to diagonal entries (see above).

  • diag_shift – Constant shift \(\epsilon_2\) added to diagonal entries (see above).

  • chunk_size – If supplied, overrides the chunk size of the variational state (useful for models where the backward pass requires more memory than the forward pass).

Return type:

QGTJacobianPyTreeT