NetKet 3.9.2__tar.gz → 3.10__tar.gz

{netket-3.9.2 → netket-3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: NetKet
-Version: 3.9.2
+Version: 3.10
 Summary: Netket : Machine Learning toolbox for many-body quantum systems.
 Project-URL: homepage, https://www.netket.org
 Project-URL: documentation, https://netket.readthedocs.io/en/latest/#
@@ -17,23 +17,26 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.8
-Requires-Dist: flax<0.8,>=0.6.5
-Requires-Dist: igraph<0.11.0,>=0.10.0
+Requires-Python: >=3.9
+Requires-Dist: beartype>=0.16.2
+Requires-Dist: flax<0.8,>=0.7.1
+Requires-Dist: igraph<0.12.0,>=0.10.0
 Requires-Dist: jax<0.5,>=0.4.3
 Requires-Dist: jaxlib<0.5,>=0.4.3
 Requires-Dist: numba4jax<0.1,>=0.0.10.post1
-Requires-Dist: numba<0.58,>=0.52
+Requires-Dist: numba<0.59,>=0.52
 Requires-Dist: numpy~=1.20
 Requires-Dist: optax<0.2,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
-Requires-Dist: plum-dispatch<3,>=1.5.1
+Requires-Dist: rich>=10.0
 Requires-Dist: scipy<2,>=1.5.3
 Requires-Dist: tqdm<5,>=4.60
+Requires-Dist: typing-extensions; python_version <= '3.10'
 Provides-Extra: dev
-Requires-Dist: black==22.10.0; extra == 'dev'
+Requires-Dist: black==23.9.1; extra == 'dev'
 Requires-Dist: build; extra == 'dev'
 Requires-Dist: coverage>=5; extra == 'dev'
+Requires-Dist: ipython; extra == 'dev'
 Requires-Dist: matplotlib>=3; extra == 'dev'
 Requires-Dist: networkx<4,>=2.4; extra == 'dev'
 Requires-Dist: pre-commit>=2.7; extra == 'dev'
@@ -41,25 +44,28 @@ Requires-Dist: pytest-cov>=2.10.1; extra == 'dev'
 Requires-Dist: pytest-json-report>=1.3; extra == 'dev'
 Requires-Dist: pytest-xdist[psutil]>=2; extra == 'dev'
 Requires-Dist: pytest>=6; extra == 'dev'
-Requires-Dist: ruff==0.0.278; extra == 'dev'
+Requires-Dist: ruff==0.1.0; extra == 'dev'
 Requires-Dist: wheel; extra == 'dev'
 Provides-Extra: docs
-Requires-Dist: myst-nb~=0.17.1; extra == 'docs'
+Requires-Dist: myst-nb~=1.0.0rc0; extra == 'docs'
 Requires-Dist: myst-parser<2.1.0,>=0.18.1; extra == 'docs'
-Requires-Dist: nbsphinx~=0.9.1; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints~=1.22; extra == 'docs'
 Requires-Dist: sphinx-book-theme~=1.0.1; extra == 'docs'
-Requires-Dist: sphinx<7.1,>=5.3; extra == 'docs'
+Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
+Requires-Dist: sphinx<7.3,>=6.2; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
-Requires-Dist: sphinxcontrib-napoleon==0.7; extra == 'docs'
 Provides-Extra: extra
 Requires-Dist: h5py>=3.7.0; extra == 'extra'
 Requires-Dist: openfermion>=1.0.0; extra == 'extra'
+Requires-Dist: qutip>=4; extra == 'extra'
 Requires-Dist: tensorboardx>=2.0.0; extra == 'extra'
 Provides-Extra: mpi
-Requires-Dist: mpi4jax<0.4,>=0.3.9; extra == 'mpi'
+Requires-Dist: mpi4jax<0.5,>=0.3.9; extra == 'mpi'
 Requires-Dist: mpi4py<4,>=3.0.1; extra == 'mpi'
+Provides-Extra: pyscf
+Requires-Dist: pyscf>=2.0; extra == 'pyscf'
+Requires-Dist: sparse>=0.12.0; extra == 'pyscf'
 Description-Content-Type: text/markdown
 
 <div align="center">
@@ -158,7 +164,7 @@ If you want MPI support, please follow the discussion in [mpi4jax](https://githu
 
 ## Getting Started
 
-To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.htmls), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
+To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.html), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
 There are also many example scripts that you can download, run and edit that showcase some use-cases of NetKet, although they are not commented.
 
 If you want to get in touch with us, feel free to open an issue or a discussion here on GitHub, or to join the MLQuantum slack group where several people involved with NetKet hang out. To join the slack channel just accept [this invitation](https://join.slack.com/t/mlquantum/shared_invite/zt-19wibmfdv-LLRI6i43wrLev6oQX0OfOw)

{netket-3.9.2 → netket-3.10}/README.md

@@ -94,7 +94,7 @@ If you want MPI support, please follow the discussion in [mpi4jax](https://githu
 
 ## Getting Started
 
-To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.htmls), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
+To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.html), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
 There are also many example scripts that you can download, run and edit that showcase some use-cases of NetKet, although they are not commented.
 
 If you want to get in touch with us, feel free to open an issue or a discussion here on GitHub, or to join the MLQuantum slack group where several people involved with NetKet hang out. To join the slack channel just accept [this invitation](https://join.slack.com/t/mlquantum/shared_invite/zt-19wibmfdv-LLRI6i43wrLev6oQX0OfOw)

{netket-3.9.2 → netket-3.10}/netket/__init__.py

@@ -12,18 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# enable x64 on jax
-# must be done at startup.
-from jax.config import config
-
-config.update("jax_enable_x64", True)
-del config
-
 from ._version import version as __version__  # noqa: F401
 
-from . import utils
 from .utils import config
 
+from . import utils
 from . import errors
 
 __all__ = [

netket-3.10/netket/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '3.10'
+__version_tuple__ = version_tuple = (3, 10)

{netket-3.9.2 → netket-3.10}/netket/driver/abstract_variational_driver.py

@@ -22,6 +22,7 @@ import jax
 from jax.tree_util import tree_map
 
 from netket.logging import JsonLog
+from netket.operator import AbstractOperator
 from netket.utils import mpi
 
 
@@ -108,19 +109,18 @@ class AbstractVariationalDriver(abc.ABC):
     def reset(self):
         """
         Resets the driver.
-        Concrete drivers should also call super().reset() to ensure that the step
-        count is set to 0.
+
+        Subclasses should make sure to call :code:`super().reset()` to ensure
+        that the step count is set to 0.
         """
         self.state.reset()
         self._step_count = 0
-        pass
 
     @abc.abstractmethod
     def info(self, depth=0):
         """
         Returns an info string used to print information to screen about this driver.
         """
-        pass  # pragma: no cover
 
     @property
     def state(self):
@@ -156,9 +156,10 @@ class AbstractVariationalDriver(abc.ABC):
         after every `step_size` steps.
 
         Args:
-            n_iter: The total number of steps to perform.
-            step_size: The number of internal steps the simulation
-                is advanced every turn.
+            n_steps: The total number of steps to perform (this is
+                equivalent to the length of the iterator)
+            step: The number of internal steps the simulation
+                is advanced between yielding from the iterator
 
         Yields:
             int: The current step.
@@ -176,7 +177,9 @@ class AbstractVariationalDriver(abc.ABC):
         """
         Performs `steps` optimization steps.
 
-        steps: (Default=1) number of steps
+        Args:
+            steps: (Default=1) number of steps.
+
         """
         for _ in self.iter(steps):
             pass
@@ -254,7 +257,6 @@ class AbstractVariationalDriver(abc.ABC):
             first_step = True
 
             for step in self.iter(n_iter, step_size):
-
                 log_data = self.estimate(obs)
                 self._log_additional_data(log_data, step)
 
@@ -306,7 +308,14 @@ class AbstractVariationalDriver(abc.ABC):
             A pytree of the same structure as the input, containing MCMC statistics
             for the corresponding operators as leaves.
         """
-        return tree_map(self._estimate_stats, observables)
+
+        # Do not unpack operators, even if they are pytrees!
+        # this is necessary to support jax operators.
+        return tree_map(
+            self._estimate_stats,
+            observables,
+            is_leaf=lambda x: isinstance(x, AbstractOperator),
+        )
 
     def update_parameters(self, dp):
         """
@@ -333,7 +342,6 @@ class AbstractVariationalDriver(abc.ABC):
         Returns:
             Nothing. The log dictionary should be modified in place.
         """
-        pass  # pragma: no cover
 
 
 @partial(jax.jit, static_argnums=0)

{netket-3.9.2 → netket-3.10}/netket/driver/steady_state.py

@@ -119,8 +119,8 @@ class SteadyState(AbstractVariationalDriver):
         argument is the step, used to change some parameters along the
         optimisation.
 
-        Often, this is taken to be :func:`nk.optimizer.SR`. If it is set to
-        `None`, then the identity is used.
+        Often, this is taken to be :func:`~netket.optimizer.SR`. If it is
+        set to `None`, then the identity is used.
         """
         return self._preconditioner
 
@@ -158,4 +158,4 @@ class SteadyState(AbstractVariationalDriver):
                 ("SR solver   ", self.sr),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])

{netket-3.9.2 → netket-3.10}/netket/driver/vmc.py

@@ -103,8 +103,8 @@ class VMC(AbstractVariationalDriver):
         argument is the step, used to change some parameters along the
         optimisation.
 
-        Often, this is taken to be :func:`nk.optimizer.SR`. If it is set to
-        `None`, then the identity is used.
+        Often, this is taken to be :func:`~netket.optimizer.SR`. If it is
+        set to `None`, then the identity is used.
         """
         return self._preconditioner
 
@@ -169,4 +169,4 @@ class VMC(AbstractVariationalDriver):
                 ("State          ", self.state),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])

{netket-3.9.2 → netket-3.10}/netket/errors.py

@@ -599,6 +599,29 @@ class RealQGTComplexDomainError(Exception):
         )
 
 
+class UnoptimalSRtWarning(NetketWarning):
+    """
+    SRt should be used when the number of parameters exceed the number of samples.
+    If this is not the case, employing `netket.driver.VMC` with the `nk.optimizer.SR` preconditioner
+    is a more efficient option while maintaining the same parameter dynamics.
+
+    .. note::
+        a  detailed discussion can be found in the documentation of
+        :func:`netket.experimental.driver.VMC_SRt`).
+
+    """
+
+    def __init__(self, n_parameters, n_samples):
+        super().__init__(
+            f"""
+            You are in the case n_samples > num_params ({n_samples} > {n_parameters}),
+            for which the `VMC_SRt` is not optimal. Consider using `netket.driver.VMC`
+            with the preconditioner `nk.optimizer.SR` to achieve the same parameter dynamics,
+            but with improved speed.
+            """
+        )
+
+
 #################################################
 # Functions to throw errors                     #
 #################################################

{netket-3.9.2 → netket-3.10}/netket/exact.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 import numpy as _np
 from scipy.sparse.linalg import bicgstab as _bicgstab
 from scipy.sparse.linalg import LinearOperator as _LinearOperator
@@ -26,7 +28,7 @@ def lanczos_ed(
     k: int = 1,
     compute_eigenvectors: bool = False,
     matrix_free: bool = False,
-    scipy_args: dict = None,
+    scipy_args: Optional[dict] = None,
 ):
     r"""Computes `first_n` smallest eigenvalues and, optionally, eigenvectors
     of a Hermitian operator using :meth:`scipy.sparse.linalg.eigsh`.

{netket-3.9.2 → netket-3.10}/netket/experimental/__init__.py

@@ -18,9 +18,11 @@ __all__ = [
     "sampler",
     "vqs",
     "TDVP",
+    "models",
     "hilbert",
     "operator",
     "logging",
+    "observable",
 ]
 
 from . import hilbert
@@ -28,9 +30,11 @@ from . import operator
 from . import driver
 from . import dynamics
 from . import sampler
+from . import models
 from . import vqs
 from . import logging
 from . import qsr
+from . import observable
 
 from .driver import TDVP
 from .qsr import QSR

{netket-3.9.2 → netket-3.10}/netket/experimental/driver/__init__.py

@@ -14,7 +14,7 @@
 
 from .tdvp import TDVP
 from .tdvp_schmitt import TDVPSchmitt
-
+from .vmc_srt import VMC_SRt
 from netket.utils import _hide_submodules
 
 _hide_submodules(__name__)

{netket-3.9.2 → netket-3.10}/netket/experimental/driver/tdvp.py

@@ -135,7 +135,7 @@ class TDVP(TDVPBaseDriver):
                 ("state         ", self.state),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])
 
 
 @odefun.dispatch
@@ -179,7 +179,6 @@ def odefun_tdvp(  # noqa: F811
 
 @partial(jax.jit, static_argnums=(3, 4))
 def _map_parameters(forces, parameters, loss_grad_factor, propagation_type, state_T):
-
     forces = jax.tree_map(
         lambda x, target: loss_grad_factor * x,
         forces,

{netket-3.9.2 → netket-3.10}/netket/experimental/driver/tdvp_common.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Optional, Union
+from collections.abc import Sequence
 import warnings
 
 import jax
@@ -205,6 +206,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
                 at which this method will stop and yield. By default, a stop is performed
                 after each time step (at potentially varying step size if an adaptive
                 integrator is used).
+
         Yields:
             The current step count.
         """
@@ -214,7 +216,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         self,
         T: float,
         tstops: Optional[Sequence[float]] = None,
-        callback: Callable = None,
+        callback: Optional[Callable] = None,
     ):
         """
         Implementation of :code:`iter`. This method accepts and additional `callback` object, which

{netket-3.9.2 → netket-3.10}/netket/experimental/driver/tdvp_schmitt.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Union
+from typing import Callable, Union, Optional
 
 from functools import partial
 
@@ -108,9 +108,9 @@ class TDVPSchmitt(TDVPBaseDriver):
         *,
         t0: float = 0.0,
         propagation_type="real",
-        holomorphic: bool = None,
+        holomorphic: Optional[bool] = None,
         diag_shift: float = 0.0,
-        diag_scale: float = None,
+        diag_scale: Optional[float] = None,
         error_norm: Union[str, Callable] = "qgt",
         rcond: float = 1e-14,
         rcond_smooth: float = 1e-8,
@@ -192,7 +192,7 @@ class TDVPSchmitt(TDVPBaseDriver):
                 ("state         ", self.state),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])
 
 
 # Copyright notice:
@@ -295,7 +295,6 @@ def odefun_schmitt(state: MCState, self: TDVPSchmitt, t, w, *, stage=0):  # noqa
 
 @partial(jax.jit, static_argnums=(3, 4))
 def _map_parameters(forces, parameters, loss_grad_factor, propagation_type, state_T):
-
     forces = jax.tree_map(
         lambda x, target: loss_grad_factor * x,
         forces,

netket-3.10/netket/experimental/driver/vmc_srt.py

@@ -0,0 +1,248 @@
+from functools import partial
+from typing import Callable, Optional
+
+import jax
+import jax.numpy as jnp
+import jax.scipy as jsp
+from netket.driver.vmc import VMC
+
+import netket.jax as nkjax
+import netket.stats as nkstats
+from netket.operator import AbstractOperator
+from netket.errors import UnoptimalSRtWarning
+import warnings
+
+from netket.vqs import MCState
+from netket.utils import mpi
+from netket.utils.types import ScalarOrSchedule
+
+from jax.flatten_util import ravel_pytree
+
+
+@partial(jax.jit, static_argnames=("mode", "solver_fn"))
+def SRt(O_L, local_energies, diag_shift, *, mode, solver_fn, e_mean=None):
+    """
+    For more details, see `https://arxiv.org/abs/2310.05715'. In particular,
+    the following parallel implementation is described in Appendix "Distributed SR computation".
+    """
+    N_params = O_L.shape[-1]
+    N_mc = O_L.shape[0] * mpi.n_nodes
+
+    local_energies = local_energies.flatten()
+
+    if e_mean is None:
+        e_mean = mpi.mean(local_energies)
+    de = jnp.conj(local_energies - e_mean).squeeze()
+
+    # * in this case O_L should be padded with zeros
+    assert (N_params % mpi.n_nodes) == 0
+
+    O_L = O_L / N_mc**0.5
+    dv = -2.0 * de / N_mc**0.5
+
+    if mode == "complex":
+        O_L = jnp.concatenate((O_L[:, 0], O_L[:, 1]), axis=0)
+        dv = jnp.concatenate((jnp.real(dv), -jnp.imag(dv)), axis=-1)
+    elif mode == "real":
+        dv = dv.real
+    else:
+        raise NotImplementedError()
+
+    # twons, (np, n_nodes) -> twons, np, n_nodes
+    O_LT = O_L.reshape(O_L.shape[0], -1, mpi.n_nodes)
+    # twons, np, n_nodes -> n_nodes, twons, np
+    O_LT = jnp.moveaxis(O_LT, -1, 0)
+
+    dv, token = mpi.mpi_gather_jax(dv)
+    dv = dv.reshape(-1, *dv.shape[2:])
+    O_LT, token = mpi.mpi_alltoall_jax(O_LT, token=token)
+
+    # proc, twons, np -> (proc, twons) np
+    O_LT = O_LT.reshape(-1, O_LT.shape[-1])
+
+    matrix, token = mpi.mpi_reduce_sum_jax(O_LT @ O_LT.T, token=token)
+    matrix_side = matrix.shape[-1]  # * it can be Ns or 2*Ns, depending on mode
+
+    if mpi.rank == 0:
+        matrix = matrix + diag_shift * jnp.eye(
+            matrix_side
+        )  # * shift diagonal regularization
+        aus_vector = solver_fn(matrix, dv)
+        aus_vector = aus_vector.reshape(mpi.n_nodes, -1)
+        aus_vector, token = mpi.mpi_scatter_jax(aus_vector, token=token)
+    else:
+        shape = jnp.zeros((int(matrix_side / mpi.n_nodes),), dtype=jnp.float64)
+        aus_vector, token = mpi.mpi_scatter_jax(shape, token=token)
+
+    updates = O_L.T @ aus_vector
+    updates, token = mpi.mpi_allreduce_sum_jax(updates, token=token)
+
+    return -updates
+
+
+inv_default_solver = lambda A, b: jnp.linalg.inv(A) @ b
+linear_solver = lambda A, b: jsp.linalg.solve(A, b, assume_a="pos")
+
+
+@jax.jit
+def _flatten_samples(x):
+    return x.reshape(-1, x.shape[-1])
+
+
+class VMC_SRt(VMC):
+    r"""
+    Energy minimization using Variational Monte Carlo (VMC) and the kernel
+    formulation of Stochastic Reconfiguration (SR). This approach lead to
+    *exactly* the same parameter updates of the standard SR with a
+    diagonal shift regularization. For this reason, it is equivalent to the standard
+    nk.driver.VMC with the preconditioner nk.optimizer.SR(solver=netket.optimizer.solver.solvers.solve)).
+    In the kernel SR framework, the updates of the parameters can be written as:
+
+    .. math::
+        \delta \theta = \tau X(X^TX + \lambda \mathbb{I}_{2M})^{-1} f,
+
+    where :math:`X \in R^{P \times 2M}` is the concatenation of the real and imaginary part
+    of the centered Jacobian, with P the number of parameters and M the number of samples.
+    The vector f is the concatenation of the real and imaginary part of the centered local
+    energy. Note that, to compute the updates, it is sufficient to invert an :math:`M\times M` matrix
+    instead of a :math:`P\times P` one. As a consequence, this formulation is useful
+    in the typical deep learning regime where :math:`P \gg M`.
+
+    See `R.Rende, L.L.Viteritti, L.Bardone, F.Becca and S.Goldt <https://arxiv.org/abs/2310.05715>`_
+    for a detailed description of the derivation. A similar result can be obtained by minimizing the
+    Fubini-Study distance with a specific constrain, see `A.Chen and M.Heyl <https://arxiv.org/abs/2302.01941>`_
+    for details.
+    """
+
+    def __init__(
+        self,
+        hamiltonian: AbstractOperator,
+        optimizer,
+        *,
+        diag_shift: ScalarOrSchedule,
+        linear_solver_fn: Callable[[jax.Array, jax.Array], jax.Array] = linear_solver,
+        jacobian_mode: Optional[str] = None,
+        variational_state: MCState = None,
+    ):
+        """
+        Initializes the driver class.
+
+        Args:
+            hamiltonian: The Hamiltonian of the system.
+            optimizer: Determines how optimization steps are performed given the
+                    bare energy gradient.
+            diag_shift: The diagonal shift of the stochastic reconfiguration matrix.
+                        Typical values are 1e-4 ÷ 1e-3. Can also be an optax schedule.
+            hamiltonian: The Hamiltonian of the system.
+            linear_solver_fn: Callable to solve the linear problem associated to the
+                              updates of the parameters
+            jacobian_mode: The mode used to compute the jacobian of the variational state. Can be `'real'`
+                    or `'complex'` (defaults to the dtype of the output of the model).
+            variational_state: The :class:`netket.vqs.MCState` to be optimised. Other
+                variational states are not supported.
+        """
+        super().__init__(hamiltonian, optimizer, variational_state=variational_state)
+
+        self._ham = hamiltonian.collect()  # type: AbstractOperator
+        self.diag_shift = diag_shift
+        self.jacobian_mode = jacobian_mode
+        self._linear_solver_fn = linear_solver_fn
+
+        _, unravel_params_fn = ravel_pytree(self.state.parameters)
+        self._unravel_params_fn = jax.jit(unravel_params_fn)
+
+        if self.state.n_parameters % mpi.n_nodes != 0:
+            raise NotImplementedError(
+                f"""
+                VMC_SRt requires a network with a number of parameters
+                multiple of the number of MPI devices/ranks in use.
+
+                You have a network with {self.state.n_parameters}, but
+                there are {mpi.n_nodes} MPI ranks in use.
+
+                To fix this, either add some 'fake' parameters to your
+                network, or change the number of MPI nodes, or contribute
+                some padding logic to NetKet!
+                """
+            )
+
+        if self.state.n_parameters < self.state.n_samples:
+            warnings.warn(
+                UnoptimalSRtWarning(self.state.n_parameters, self.state.n_samples),
+                UserWarning,
+                stacklevel=2,
+            )
+
+    @property
+    def jacobian_mode(self) -> str:
+        """
+        The mode used to compute the jacobian of the variational state. Can be `'real'`
+        or `'complex'`.
+
+        Real mode truncates imaginary part of the wavefunction, while `complex` does not.
+        This internally uses :func:`netket.jax.jacobian`. See that function for a more
+        complete documentation.
+        """
+        return self._jacobian_mode
+
+    @jacobian_mode.setter
+    def jacobian_mode(self, mode: Optional[str]):
+        if mode is None:
+            mode = nkjax.jacobian_default_mode(
+                self.state._apply_fun,
+                self.state.parameters,
+                self.state.model_state,
+                self.state.samples,
+                warn=False,
+            )
+
+        if mode not in ["complex", "real"]:
+            raise ValueError(
+                "`jacobian_mode` only supports 'real' for real-valued wavefunctions and"
+                "'complex'.\n\n"
+                "`holomorphic` is not yet supported, but could be contributed in the future."
+            )
+        self._jacobian_mode = mode
+
+    def _forward_and_backward(self):
+        """
+        Performs a number of VMC optimization steps.
+
+        Args:
+            n_steps (int): Number of steps to perform.
+        """
+
+        self.state.reset()
+
+        # Compute the local energy estimator and average Energy
+        local_energies = self.state.local_estimators(self._ham)
+
+        self._loss_stats = nkstats.statistics(local_energies)
+
+        samples = _flatten_samples(self.state.samples)
+        jacobians = nkjax.jacobian(
+            self.state._apply_fun,
+            self.state.parameters,
+            samples,
+            self.state.model_state,
+            mode=self.jacobian_mode,
+            dense=True,
+            center=True,
+        )  # * jaxcobians is centered
+
+        diag_shift = self.diag_shift
+        if callable(self.diag_shift):
+            diag_shift = diag_shift(self.step_count)
+
+        updates = SRt(
+            jacobians,
+            local_energies,
+            diag_shift,
+            mode=self.jacobian_mode,
+            solver_fn=self._linear_solver_fn,
+            e_mean=self._loss_stats.Mean,
+        )
+
+        self._dp = self._unravel_params_fn(updates)
+
+        return self._dp