NetKet 3.22.4__tar.gz → 3.23.0__tar.gz

{netket-3.22.4 → netket-3.23.0}/.gitignore

@@ -44,6 +44,8 @@ uv.lock
 # Logs and databases #
 ######################
 *.log
+# per-rank output of djaxrun
+djaxrunlog-p*
 *.sql
 *.sqlite
 *.json

{netket-3.22.4 → netket-3.23.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: NetKet
-Version: 3.22.4
+Version: 3.23.0
 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/#
@@ -62,11 +62,11 @@ Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
 Provides-Extra: extra
 Requires-Dist: h5py>=3.7.0; extra == 'extra'
-Requires-Dist: jaxmg<1.0.0,>=0.0.9; (sys_platform == 'linux') and extra == 'extra'
 Requires-Dist: mlflow>=2.0.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'
+Requires-Dist: wandb>=0.19.0; extra == 'extra'
 Provides-Extra: pyscf
 Requires-Dist: pyscf>=2.0; extra == 'pyscf'
 Description-Content-Type: text/markdown

{netket-3.22.4 → netket-3.23.0}/netket/_src/callbacks/base.py

@@ -25,6 +25,16 @@ class StopRun(Exception):
     pass
 
 
+class RunFinished(StopRun):
+    """
+    Exception to be raised by callbacks to end a run that has reached its end,
+    for example the final time of a time evolution.
+
+    It works like :class:`~netket.callbacks.StopRun`, but the driver does not
+    print that the run is stopping early.
+    """
+
+
 #: ``callback_order`` for callbacks that may raise :class:`StopRun`. They run
 #: *last* (after observables at 0 and loggers at 10) so that, on the step a
 #: callback stops the run, no later callback in the hook is skipped — a skipped

{netket-3.22.4 → netket-3.23.0}/netket/_src/callbacks/convergence_stopping.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import deque
-
 import numpy as np
+from flax import serialization
 
 from netket.utils import struct
 
@@ -41,8 +40,11 @@ class ConvergenceStopping(AbstractCallback, mutable=True):
     """The loss must be consistently below this value for this number of
     iterations in order to stop the optimisation."""
 
-    _loss_window: deque = struct.field(pytree_node=False, serialize=False)
-    _patience_counter: int = struct.field(pytree_node=False, serialize=False, default=0)
+    # Saved, so that a run resumed from a checkpoint keeps its history. The window
+    # has a fixed length: only its last `_n_losses` entries are filled.
+    _loss_window: tuple[float, ...] = struct.field(pytree_node=False, serialize=True)
+    _n_losses: int = struct.field(pytree_node=False, serialize=True, default=0)
+    _patience_counter: int = struct.field(pytree_node=False, serialize=True, default=0)
 
     def __init__(
         self,
@@ -73,19 +75,32 @@ class ConvergenceStopping(AbstractCallback, mutable=True):
         self.smoothing_window = smoothing_window
         self.patience = patience
 
-        self._loss_window = deque([], maxlen=smoothing_window)
+        self._loss_window = (0.0,) * smoothing_window
+        self._n_losses = 0
         self._patience_counter = 0
 
+    def __process_deserialization_state__(self, state):
+        # If the file has no window (older versions) or one of a different size,
+        # start the window over.
+        if len(state.get("_loss_window", ())) != self.smoothing_window:
+            state = {
+                k: v for k, v in state.items() if k not in ("_loss_window", "_n_losses")
+            }
+        names = ("_loss_window", "_n_losses", "_patience_counter")
+        defaults = {n: serialization.to_state_dict(getattr(self, n)) for n in names}
+        return {**defaults, **state}
+
     @property
     def callback_order(self) -> int:
         # Run last, so raising StopRun never skips a later callback's collective.
         return STOPPING_CALLBACK_ORDER
 
     def on_step_end(self, step, log_data, driver):
-        loss = np.asarray(np.real(getattr(log_data[driver._loss_name], self.monitor)))
+        loss = np.real(getattr(log_data[driver._loss_name], self.monitor))
 
-        self._loss_window.append(loss)
-        loss_smooth = np.mean(self._loss_window)
+        self._loss_window = (*self._loss_window[1:], float(np.mean(loss)))
+        self._n_losses = min(self._n_losses + 1, self.smoothing_window)
+        loss_smooth = np.mean(self._loss_window[-self._n_losses :])
 
         if loss_smooth <= self.target:
             self._patience_counter += 1

{netket-3.22.4 → netket-3.23.0}/netket/_src/callbacks/early_stopping.py

@@ -45,12 +45,14 @@ class EarlyStopping(AbstractCallback, mutable=True):
     start_from_step: int = struct.field(pytree_node=False)
     """Number of steps to wait before the callback has any effect."""
 
-    _best_val: float = struct.field(pytree_node=False, serialize=False, default=np.inf)
+    # Saved, so that a run resumed from a checkpoint keeps counting its patience
+    # instead of starting over.
+    _best_val: float = struct.field(pytree_node=False, serialize=True, default=np.inf)
     """Best value of the loss observed up to this iteration."""
-    _best_iter: int = struct.field(pytree_node=False, serialize=False, default=0)
+    _best_iter: int = struct.field(pytree_node=False, serialize=True, default=0)
     """Iteration at which the `_best_val` was observed."""
     _best_patience_counter: int = struct.field(
-        pytree_node=False, serialize=False, default=0
+        pytree_node=False, serialize=True, default=0
     )
     """Stores how many steps have elapsed without an improvement."""
 
@@ -91,6 +93,14 @@ class EarlyStopping(AbstractCallback, mutable=True):
         self._best_iter = 0
         self._best_patience_counter = 0
 
+    # TODO: backward compatibility for checkpoints saved before September 2026.
+    # Remove in 2027.
+    def __process_deserialization_state__(self, state):
+        # Files saved by older versions do not contain the best-so-far state.
+        names = ("_best_val", "_best_iter", "_best_patience_counter")
+        defaults = {name: self._pytree__fields[name].default for name in names}
+        return {**defaults, **state}
+
     @property
     def callback_order(self) -> int:
         # Run last, so raising StopRun never skips a later callback's collective.
@@ -104,7 +114,7 @@ class EarlyStopping(AbstractCallback, mutable=True):
 
         self._best_patience_counter += 1
         if self._is_improvement(loss, self._best_val):
-            self._best_val = loss
+            self._best_val = float(loss)
             self._best_iter = step
 
             if self.baseline is None:

{netket-3.22.4 → netket-3.23.0}/netket/_src/callbacks/invalid_loss_stopping.py

@@ -33,7 +33,9 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
     patience: int | float = struct.field(pytree_node=False)
     """Number of epochs with invalid loss after which training will be stopped."""
 
-    _last_valid_iter: int = struct.field(pytree_node=False, serialize=False, default=0)
+    # Saved, so that a run resumed from a checkpoint does not stop at the first
+    # invalid step, as if the loss had been invalid since step 0.
+    _last_valid_iter: int = struct.field(pytree_node=False, serialize=True, default=0)
     """Last valid iteration, to check against patience."""
 
     def __init__(self, monitor: str = "mean", patience: int | float = 0):
@@ -54,6 +56,12 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
         self.patience = patience
         self._last_valid_iter = 0
 
+    # TODO: backward compatibility for checkpoints saved before September 2026.
+    # Remove in 2027.
+    def __process_deserialization_state__(self, state):
+        # Files saved by older versions do not contain it.
+        return {"_last_valid_iter": self._last_valid_iter, **state}
+
     @property
     def callback_order(self) -> int:
         # Run last, so raising StopRun never skips a later callback's collective.
@@ -67,7 +75,9 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
         if driver._loss_stats is not None:
             loss = np.real(getattr(driver._loss_stats, self.monitor))
 
-            if not np.isfinite(loss):
+            # `loss` may be vector-valued (e.g. a foundation ReplicaStats has one
+            # mean per anchor); stop if ANY component is non-finite.
+            if not np.all(np.isfinite(loss)):
                 if driver.step_count - self._last_valid_iter >= self.patience:
                     raise StopRun(
                         f"InvalidLossStopping: loss is not finite ({loss}) "

{netket-3.22.4 → netket-3.23.0}/netket/_src/callbacks/progressbar.py

@@ -10,28 +10,39 @@ from netket.utils import struct
 from netket._src.callbacks.base import AbstractCallback
 
 
+def _elapsed_time(driver, t_start: float, T_final: float) -> float:
+    """Time elapsed since `t_start`, at most `T_final`."""
+    return min(float(driver.t) - t_start, T_final)
+
+
 class TimeProgressBarCallback(AbstractCallback, mutable=True):
     """
     Progress bar displaying physical time for dynamics drivers.
 
-    Tracks ``driver.t ∈ [0, T_final]`` instead of step count.
+    Tracks the time elapsed since the start of the run, ``driver.t - t_start ∈
+    [0, T_final]``, instead of step count. ``t_start`` is ``driver._t_start`` if
+    the driver sets it (as the TDVP drivers do), and 0 otherwise.
     Used automatically by :class:`~netket.driver.AbstractDynamicsDriver`
     when :meth:`run` is called with a float time argument.
     """
 
     _T_final: float = struct.field(pytree_node=False, serialize=False)
-    _last_t: float = struct.field(pytree_node=False, serialize=False)
+    _t_start: float = struct.field(pytree_node=False, serialize=False)
     _pbar: Any = struct.field(pytree_node=False, serialize=False)
     _leave: bool = struct.field(pytree_node=False, serialize=False)
 
     def __init__(self, T_final: float, leave: bool = True):
         self._T_final = T_final
-        self._last_t = 0.0
+        self._t_start = 0.0
         self._pbar = None
         self._leave = leave
 
     def on_run_start(self, step, driver):
-        self._last_t = float(driver.t)
+        # Count from the time the run started at, even if another callback
+        # changes the time (e.g. when resuming from a checkpoint).
+        t_start = getattr(driver, "_t_start", None)
+        self._t_start = 0.0 if t_start is None else float(t_start)
+
         dt = float(driver.dt)
         n_dec = max(4, math.ceil(-math.log10(dt))) if dt > 0 else 4
         bar_format = (
@@ -46,7 +57,7 @@ class TimeProgressBarCallback(AbstractCallback, mutable=True):
             leave=self._leave,
             bar_format=bar_format,
         )
-        self._pbar.update(float(driver.t))
+        self._pbar.update(_elapsed_time(driver, self._t_start, self._T_final))
         self._pbar.unpause()
         self._pbar.refresh()
 
@@ -57,16 +68,16 @@ class TimeProgressBarCallback(AbstractCallback, mutable=True):
             )
 
     def on_step_end(self, step, log_data, driver):
-        t = float(driver.t)
-        self._pbar.update(t - self._last_t)
-        self._last_t = t
+        self._pbar.update(
+            _elapsed_time(driver, self._t_start, self._T_final) - self._pbar.n
+        )
         self._pbar.refresh()
 
     def on_run_end(self, step, driver):
         if self._pbar is not None:
-            t = float(driver.t)
-            self._pbar.update(t - self._last_t)
-            self._last_t = t
+            self._pbar.update(
+                _elapsed_time(driver, self._t_start, self._T_final) - self._pbar.n
+            )
             self._pbar.close()
             self._pbar = None
 
@@ -80,7 +91,7 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
     _leave: bool = struct.field(pytree_node=False, serialize=False)
 
     _n_steps: int = struct.field(pytree_node=False, serialize=False)
-    _last_step: int = struct.field(pytree_node=False, serialize=False)
+    _start_step: int = struct.field(pytree_node=False, serialize=False)
 
     _pbar: Any = struct.field(pytree_node=False, serialize=False)
 
@@ -90,7 +101,11 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
         self._leave = leave
 
     def on_run_start(self, step, driver):
-        self._last_step = driver.step_count
+        # Count from the step the run started at, even if another callback
+        # changes the step count (e.g. when resuming from a checkpoint).
+        self._start_step = getattr(driver, "_start_step", None)
+        if self._start_step is None:
+            self._start_step = driver.step_count
 
         self._pbar = tqdm(
             total=self._n_steps,
@@ -98,7 +113,7 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
             dynamic_ncols=True,
             leave=self._leave,
         )
-        self._pbar.update(driver.step_count - self._last_step)
+        self._pbar.update(driver.step_count - self._start_step)
         self._pbar.unpause()
         self._pbar.refresh()
 
@@ -110,12 +125,12 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
             )
 
     def on_step_end(self, step, log_data, driver):
-        step_value = driver.step_count - self._last_step
+        step_value = driver.step_count - self._start_step
         self._pbar.update(step_value - self._pbar.n)
         self._pbar.refresh()
 
     def on_run_end(self, step, driver):
-        step_value = driver.step_count - self._last_step
+        step_value = driver.step_count - self._start_step
         if self._pbar is not None:
             self._pbar.update(step_value - self._pbar.n)
             self._pbar.refresh()

{netket-3.22.4 → netket-3.23.0}/netket/_src/callbacks/save_state.py

@@ -40,7 +40,7 @@ class SaveVariationalState(AbstractCallback, mutable=True):
         >>> driver.run(  # doctest: +SKIP
         ...     n_iter=50,
         ...     out="test",
-        ...     callback=nk.callbacks.SaveVariationalStateCallback(path="optimization", interval=10),
+        ...     callback=nk.logging.SaveVariationalState(path="optimization", interval=10),
         ... )
         >>> nqxpack.load("optimization/state_00010.nk")  # doctest: +SKIP
 

{netket-3.22.4 → netket-3.23.0}/netket/_src/driver/abstract_dynamics_driver.py

@@ -166,8 +166,7 @@ class AbstractDynamicsDriver(AbstractDriver):
                 callbacks fire normally.
             out: Logger or iterable of loggers for output.
             obs: Observables to compute at each logging step.
-            max_steps: Safety cap on iterations when ``T_or_n`` is a float,
-                to prevent infinite loops if ``dt`` is zero.
+            n_iter: Number of iterations to perform. Mutually exclusive with ``T``.
             show_progress: Show a progress bar (default True).
             callback: User callbacks.
             timeit: If True, print timing information after the run.

{netket-3.22.4 → netket-3.23.0}/netket/_src/driver/abstract_optimization_driver.py

@@ -31,9 +31,10 @@ class AbstractOptimizationDriver(AbstractDriver):
     Adds an `optax <https://optax.readthedocs.io/en/latest/>`_ optimizer and
     implements :meth:`update_parameters` via :func:`apply_gradient`.
 
-    Subclass this when implementing a new optimization driver. The public
-    name :class:`AbstractVariationalDriver` is currently kept as a compatibility
-    alias for this class.
+    Subclass this when implementing a new optimization driver, following the
+    guidelines in :class:`~netket.driver.AbstractDriver` (including how to support
+    checkpointing). The public name :class:`AbstractVariationalDriver` is
+    currently kept as a compatibility alias for this class.
     """
 
     _optimizer: Optimizer = struct.field(pytree_node=False, serialize=False)

{netket-3.22.4 → netket-3.23.0}/netket/_src/driver/abstract_variational_driver.py

@@ -29,7 +29,7 @@ from netket.utils.iterators import to_iterable
 from netket.utils.types import PyTree
 from netket.vqs import VariationalState, FullSumState
 
-from netket._src.callbacks.base import AbstractCallback, StopRun
+from netket._src.callbacks.base import AbstractCallback, StopRun, RunFinished
 from netket._src.callbacks.observables import ObservableCallback
 from netket._src.callbacks.legacy_wrappers import (
     LegacyCallbackT,
@@ -85,13 +85,19 @@ class AbstractDriver(struct.Pytree, mutable=True):
         - :meth:`~netket.driver.AbstractDriver.reset_step` (optional):
           reset the sampler at the start of each step.
 
+        To support checkpointing, set every saved field in ``__init__``, with the
+        same structure, shapes and dtypes it will have after running, instead of
+        creating it at the first step. A checkpoint can then be restored into a
+        newly created driver. Values recomputed at every step should not be saved
+        (``serialize=False``, see :func:`netket.utils.struct.field`).
+
     """
 
     # Configuration fields, not very important
     _loss_name: str = struct.field(pytree_node=False, serialize=False)
 
     # Stuff to iterate the driver
-    _loss_stats: Any = struct.field(serialize=True)
+    _loss_stats: Any = struct.field(serialize=False, default=None)
     _step_count: int = struct.field(serialize=True)
 
     # state of the driver
@@ -104,10 +110,13 @@ class AbstractDriver(struct.Pytree, mutable=True):
     # Internal caches (those could be removed in the future?)
     _dp: PyTree = struct.field(pytree_node=True, serialize=False)
 
-    # Iterator caches
-    # _step_start: int = struct.field(pytree_node=False, serialize=False, default=None)
-    # _step_size: int = struct.field(pytree_node=False, serialize=False, default=1)
-    # _step_end: int = struct.field(pytree_node=False, serialize=False, default=None)
+    # The steps at which the current (or last) call to `run` started and stops.
+    _start_step: int | None = struct.field(
+        pytree_node=False, serialize=False, default=None
+    )
+    _target_step: int | None = struct.field(
+        pytree_node=False, serialize=False, default=None
+    )
     _step_attempt: int = struct.field(pytree_node=False, serialize=False, default=0)
     _timer: timing.Timer = struct.field(pytree_node=False, serialize=False)
 
@@ -395,7 +404,8 @@ class AbstractDriver(struct.Pytree, mutable=True):
         without reading the json output.
 
         Args:
-            n_iter: the total number of iterations to be performed during this run.
+            n_iter: the number of steps to run, counted from the step count when
+                `run` is called.
             out: A logger object, or an iterable of loggers, to be used to store simulation log and data.
                 If this argument is a string, it will be used as output prefix for the standard JSON logger.
             obs: An iterable containing all observables that should be computed
@@ -434,14 +444,16 @@ class AbstractDriver(struct.Pytree, mutable=True):
             show_progress=show_progress,
         )
 
-        # self._step_size = step_size
-        # self._step_start = self.step_count
-        # self._step_end = self.step_count + n_iter
+        # Fixed before `on_run_start`, so that a callback restoring `step_count`
+        # there (e.g. when resuming from a checkpoint) does not move the end of
+        # the run.
+        self._start_step = self.step_count
+        self._target_step = self.step_count + n_iter
 
         with timing.timed_scope(force=timeit) as timer:
             try:
                 callbacks.on_run_start(self.step_count, self)
-                for step in range(self.step_count, self.step_count + n_iter):
+                for step in range(self.step_count, self._target_step):
                     self._step_attempt = 0
                     step_log_data = {}
 
@@ -481,7 +493,11 @@ class AbstractDriver(struct.Pytree, mutable=True):
                 callbacks.on_run_end(self.step_count, self)
                 # If all processes did not raise StopRun, some will deadlock here.
                 # This means you have a bug somewhere in YOUR code.
-                if jax.process_count() > 1:
+                if isinstance(error, RunFinished):
+                    # The run reached its end (e.g. the final time of a time
+                    # evolution), it did not stop early.
+                    pass
+                elif jax.process_count() > 1:
                     print(
                         f"[process {jax.process_index()}/{jax.process_count()}] "
                         f"Stopping early because of : {error}"

{netket-3.22.4 → netket-3.23.0}/netket/_src/driver/infidelity_sr.py

@@ -22,6 +22,7 @@ from netket._src.driver.abstract_optimization_driver import (
 )
 from netket._src.ngd.sr_srt_common import sr, srt, get_samples_and_pdf
 from netket._src.ngd.srt_onthefly import srt_onthefly
+from netket._src.driver.vmc_sr import _init_momentum_buffer
 from netket._src.nn.apply_operator.functional import make_logpsi_op_afun
 from netket._src.callbacks.auto_chunk_size import get_forward_operator
 from netket._src.observable.infidelity.expect import get_local_estimator
@@ -109,7 +110,7 @@ class Infidelity_SR(AbstractOptimizationDriver):
 
     # Serialized state
     _old_updates: PyTree = None
-    info: Any | None = None
+    info: Any | None = struct.field(serialize=False, default=None)
     """
     PyTree to pass on information from the solver,e.g, the quadratic model.
     """
@@ -257,7 +258,11 @@ class Infidelity_SR(AbstractOptimizationDriver):
         _, unravel_params_fn = ravel_pytree(self.state.parameters)
         self._unravel_params_fn = jax.jit(unravel_params_fn)
 
-        self._old_updates: PyTree = None
+        # Built here rather than at the first step, so that a freshly built
+        # driver can load it back from a checkpoint.
+        self._old_updates = _init_momentum_buffer(
+            self.state.parameters, self.momentum, self._on_the_fly
+        )
 
         # PyTree to pass on information from the solver, e.g, the quadratic model
         self.info = None

{netket-3.22.4 → netket-3.23.0}/netket/_src/driver/vmc_sr.py

@@ -73,6 +73,37 @@ def VMC_SRt(
     )
 
 
+# Number of lags used for the MCMC diagnostics accumulated across iterations.
+_MCMC_DIAGNOSTICS_MAX_LAG = 32
+
+
+def _init_momentum_buffer(parameters, momentum, on_the_fly):
+    """The zeros that the SPRING solvers would otherwise create at the first step."""
+    if momentum is None:
+        return None
+    params_real, _ = nkjax.tree_to_real(parameters)
+    buffer = jax.tree_util.tree_map(jnp.zeros_like, params_real)
+    if not on_the_fly:
+        buffer, _ = ravel_pytree(buffer)
+    return buffer
+
+
+def _init_mcmc_diagnostics(state, hamiltonian, decay):
+    """The empty accumulator that online_statistics would create at the first step."""
+    # It must have the same dtype as the local energies.
+    log_psi = jax.eval_shape(
+        state._apply_fun,
+        state.variables,
+        state.hilbert.random_state(jax.random.key(0), 1),
+    )
+    return OnlineStats(
+        state.sampler.n_chains,
+        dtype=jnp.result_type(hamiltonian.dtype, log_psi.dtype),
+        decay=decay,
+        max_lag=_MCMC_DIAGNOSTICS_MAX_LAG,
+    )
+
+
 @reference(
     "Goldshlager2023Spring",
     condition="If using VMC_SR with momentum != 0",
@@ -247,7 +278,7 @@ class VMC_SR(AbstractOptimizationDriver):
 
     # Serialized state
     _old_updates: PyTree = None
-    info: Any | None = None
+    info: Any | None = struct.field(serialize=False, default=None)
     """
     PyTree to pass on information from the solver,e.g, the quadratic model.
     """
@@ -410,7 +441,11 @@ class VMC_SR(AbstractOptimizationDriver):
         _, unravel_params_fn = ravel_pytree(self.state.parameters)
         self._unravel_params_fn = jax.jit(unravel_params_fn)
 
-        self._old_updates: PyTree = None
+        # Built here rather than at the first step, so that a freshly built
+        # driver can load it back from a checkpoint.
+        self._old_updates = _init_momentum_buffer(
+            self.state.parameters, self.momentum, self._on_the_fly
+        )
 
         # PyTree to pass on information from the solver, e.g, the quadratic model
         self.info = None
@@ -429,6 +464,11 @@ class VMC_SR(AbstractOptimizationDriver):
         self._loss_stats_online = None
 
         decay = self._mcmc_convergence_diagnostics_ema_decay
+        if decay is not None:
+            # Built here so that a fresh driver can load it from a checkpoint.
+            self._loss_stats_online = _init_mcmc_diagnostics(
+                self.state, self._ham, decay
+            )
         if jax.process_index() == 0 and decay is not None:
             print(
                 f"online_statistics: "
@@ -462,11 +502,16 @@ class VMC_SR(AbstractOptimizationDriver):
             # across different MCMC iterations with an exponential moving average.
             decay = self._mcmc_convergence_diagnostics_ema_decay
             if decay is not None:
+                old_estimator = self._loss_stats_online
+                # If it is still empty, rebuild it in case the number of samples
+                # or chains changed since the driver was built.
+                if old_estimator is not None and old_estimator.n_samples == 0:
+                    old_estimator = None
                 self._loss_stats_online = nkstats.online_statistics(
                     local_energies,
-                    old_estimator=self._loss_stats_online,
+                    old_estimator=old_estimator,
                     decay=decay,
-                    max_lag=32,
+                    max_lag=_MCMC_DIAGNOSTICS_MAX_LAG,
                 )
 
         # Extract the hyperparameters which might be iteration dependent

{netket-3.22.4 → netket-3.23.0}/netket/_src/logging/hdf5_log.py

@@ -41,7 +41,13 @@ def _compute_chunk_shape(value: np.ndarray, chunk_size: int | None) -> tuple[int
 
 
 def _append_dataset(root, value, data, *, chunk_size: int | None = None):
+    # ``np.asarray`` materialises ``value`` to host -- a collective when it is a
+    # sharded array, so every rank must reach it. ``data`` is ``None`` on
+    # non-master ranks (they hold no writer): they still materialise (to join
+    # the collective) but do not write. See HDF5Log.__call__.
     value = np.asarray(value)
+    if data is None:
+        return
     if root in data:
         f_value = data[root]
         f_value.resize(f_value.shape[0] + 1, axis=0)
@@ -100,7 +106,7 @@ def tree_log(tree, root, data, *, iter=None, chunk_size: int | None = None):
 
 class HDF5Log(AbstractCallback):
     r"""
-    HDF5 Logger, that can be passed with keyword argument `logger` to Monte
+    HDF5 Logger, that can be passed with keyword argument `out` to Monte
     Carlo drivers in order to serialize the output data of the simulation.
 
     The logger has support for scalar numbers, NumPy/JAX arrays, and netket.stats.Stats objects.
@@ -117,7 +123,7 @@ class HDF5Log(AbstractCallback):
     inspecting the datasets as a dictionary, i.e. :code:`f['data/energy/Mean']`
 
     This class is a full :class:`~netket.callbacks.AbstractCallback` and can be passed
-    either as ``out=logger`` *or* inside the ``callbacks=[..., logger]`` list.
+    either as ``out=logger`` *or* as ``callback=logger``.
 
     .. tip::
         Use the ``metadata`` argument to attach a flat dict of hyper-parameters
@@ -153,7 +159,7 @@ class HDF5Log(AbstractCallback):
         >>>
         >>> import netket as nk
         >>> logger = nk.logging.HDF5Log("output")
-        >>> gs.run(n_iter=300, callbacks=[logger])
+        >>> gs.run(n_iter=300, callback=logger)
 
     .. note::
         The API of this logger is covered by our Semantic Versioning API guarantees. However, the structure of the
@@ -263,20 +269,19 @@ class HDF5Log(AbstractCallback):
         if self._writer is None:
             self._init_output_file()
 
-        if self._is_master_process:
-            tree_log(
-                log_data,
-                "data",
-                self._writer,
-                iter=step,
-                chunk_size=self._chunk_size,
-            )
+        tree_log(
+            log_data,
+            "data",
+            self._writer,
+            iter=step,
+            chunk_size=self._chunk_size,
+        )
 
         self._steps_notflushed_write += 1
         if self._steps_notflushed_write >= self._write_every:
             self.flush()
 
-    def flush(self):
+    def flush(self, variational_state=None):
         """Writes buffered data to disk."""
         if self._writer is not None:
             self._writer.flush()