NetKet 3.22.3__tar.gz → 3.22.4__tar.gz

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

@@ -21,6 +21,7 @@ pip-wheel-metadata
 .ipynb_checkpoints
 .pytest_cache
 .benchmarks
+.jax-cache
 .mypy_cache
 .virtual_documents
 .ruff_cache

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: NetKet
-Version: 3.22.3
+Version: 3.22.4
 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/#
@@ -54,7 +54,7 @@ Requires-Dist: myst-nb<1.5,>=1.0; extra == 'docs'
 Requires-Dist: myst-parser<5.2.0,>=2.0.0; extra == 'docs'
 Requires-Dist: nbconvert>=7.16.6; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints<4.0,>=1.22; extra == 'docs'
-Requires-Dist: sphinx-book-theme<1.3,>=1.1; extra == 'docs'
+Requires-Dist: sphinx-book-theme<1.5,>=1.1; extra == 'docs'
 Requires-Dist: sphinx-copybutton>=0.5.2; extra == 'docs'
 Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
 Requires-Dist: sphinx<9.2.0,>=8.2.3; extra == 'docs'
@@ -62,6 +62,7 @@ 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'

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

@@ -25,6 +25,13 @@ class StopRun(Exception):
     pass
 
 
+#: ``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
+#: collective would deadlock the run under ``jax.distributed``.
+STOPPING_CALLBACK_ORDER = 100
+
+
 class AbstractCallback(struct.Pytree, mutable=True):
     """
     Abstract base class for callbacks in advanced variational drivers.
@@ -71,7 +78,9 @@ class AbstractCallback(struct.Pytree, mutable=True):
         Lower numbers are called first, and higher numbers are called later.
 
         This can be redefined in subclasses to change the order in which callbacks are called.
-        (Default: 0, for all callbacks, 10 for loggers).
+        (Default: 0, for all callbacks, 10 for loggers, and
+        :data:`STOPPING_CALLBACK_ORDER` (100) for callbacks that may raise
+        :class:`StopRun`, so they run after everything else).
         """
         return 0
 

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

@@ -18,7 +18,11 @@ import numpy as np
 
 from netket.utils import struct
 
-from netket._src.callbacks.base import AbstractCallback, StopRun
+from netket._src.callbacks.base import (
+    AbstractCallback,
+    StopRun,
+    STOPPING_CALLBACK_ORDER,
+)
 
 
 class ConvergenceStopping(AbstractCallback, mutable=True):
@@ -72,6 +76,11 @@ class ConvergenceStopping(AbstractCallback, mutable=True):
         self._loss_window = deque([], maxlen=smoothing_window)
         self._patience_counter = 0
 
+    @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)))
 

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

@@ -16,7 +16,11 @@ import numpy as np
 
 from netket.utils import struct
 
-from netket._src.callbacks.base import AbstractCallback, StopRun
+from netket._src.callbacks.base import (
+    AbstractCallback,
+    StopRun,
+    STOPPING_CALLBACK_ORDER,
+)
 
 
 class EarlyStopping(AbstractCallback, mutable=True):
@@ -87,6 +91,11 @@ class EarlyStopping(AbstractCallback, mutable=True):
         self._best_iter = 0
         self._best_patience_counter = 0
 
+    @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):
         if step < self.start_from_step:
             return

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

@@ -16,7 +16,11 @@ import numpy as np
 
 from netket.utils import struct
 
-from netket._src.callbacks.base import AbstractCallback, StopRun
+from netket._src.callbacks.base import (
+    AbstractCallback,
+    StopRun,
+    STOPPING_CALLBACK_ORDER,
+)
 
 
 class InvalidLossStopping(AbstractCallback, mutable=True):
@@ -50,6 +54,11 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
         self.patience = patience
         self._last_valid_iter = 0
 
+    @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):
         # reset if the driver was restarted
         if driver.step_count < self._last_valid_iter:

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

@@ -1,7 +1,10 @@
+import re
+
 from typing import Any
 from pathlib import Path
 
 import jax
+import jax.experimental.multihost_utils
 
 from netket.utils import struct, timing
 from netket.utils.optional_deps import import_optional_dependency
@@ -48,6 +51,9 @@ class SaveVariationalState(AbstractCallback, mutable=True):
     _file_name_root: str = struct.field(
         pytree_node=False, serialize=False, default="state"
     )
+    _max_to_keep: int | None = struct.field(
+        pytree_node=False, serialize=False, default=None
+    )
 
     def __init__(
         self,
@@ -55,6 +61,7 @@ class SaveVariationalState(AbstractCallback, mutable=True):
         interval: float | int,
         *,
         file_name_root: str = "state",
+        max_to_keep: int | None = None,
     ):
         r"""
         Constructs the callback to save the variational state at fixed intervals.
@@ -67,12 +74,15 @@ class SaveVariationalState(AbstractCallback, mutable=True):
             path (str | Path): The path where to save the variational state.
             interval (float | int, optional): The interval at which to save the variational state
             file_name_root (str, optional): The root of the file name to save the variational state. Defaults to "state".
+            max_to_keep (int | None, optional): If set, only the most recent ``max_to_keep`` saved states are
+                kept on disk; older ones are deleted after each save. Defaults to None (keep all).
         """
         if not isinstance(path, Path):
             path = Path(path)
         self._path = path
         self._interval = interval
         self._file_name_root = file_name_root
+        self._max_to_keep = max_to_keep
 
         if jax.process_index() == 0:
             self._path.mkdir(parents=True, exist_ok=True)
@@ -82,15 +92,35 @@ class SaveVariationalState(AbstractCallback, mutable=True):
             descr="to save the variational state with SaveVariationalStateCallback.",
         )
 
+    def _save(self, step, driver):
+        path = self._path / f"{self._file_name_root}_{step:05d}.nk"
+        _nqxpack_save(driver.state, path)
+
+        if self._max_to_keep is not None:
+            if jax.process_index() == 0:
+                pattern = re.compile(rf"^{re.escape(self._file_name_root)}_(\d+)\.nk$")
+                saved = sorted(
+                    (
+                        p
+                        for p in self._path.glob(f"{self._file_name_root}_*.nk")
+                        if pattern.match(p.name)
+                    ),
+                    key=lambda p: int(pattern.match(p.name).group(1)),
+                )
+                for old in saved[: -self._max_to_keep]:
+                    old.unlink()
+            # Barrier so that other processes do not race ahead (e.g. saving the
+            # next checkpoint) while rank 0 is still pruning the directory.
+            jax.experimental.multihost_utils.sync_global_devices(
+                "SaveVariationalState._save"
+            )
+
     def before_parameter_update(self, step, log_data, driver):
         if step % self._interval == 0:
-            path = self._path / f"{self._file_name_root}_{step:05d}.nk"
-            _nqxpack_save(driver.state, path)
+            self._save(step, driver)
 
     def on_run_end(self, step, driver):
-        if step % self._interval == 0:
-            path = self._path / f"{self._file_name_root}_{step:05d}.nk"
-            _nqxpack_save(driver.state, path)
+        self._save(step, driver)
 
 
 @timing.timed

{netket-3.22.3 → netket-3.22.4}/netket/_src/callbacks/stop_at_time.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from netket._src.callbacks.base import AbstractCallback, StopRun
+from netket._src.callbacks.base import (
+    AbstractCallback,
+    StopRun,
+    STOPPING_CALLBACK_ORDER,
+)
 from netket.utils import struct
 
 
@@ -30,6 +34,11 @@ class _StopAtTime(AbstractCallback):
     def __init__(self, T_final: float):
         self.T_final = T_final
 
+    @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):
         if driver.t >= self.T_final:
             raise StopRun(f"Reached t={driver.t:.6g} >= T_final={self.T_final:.6g}")

{netket-3.22.3 → netket-3.22.4}/netket/_src/callbacks/timeout.py

@@ -16,7 +16,11 @@ import time
 
 from netket.utils import struct
 
-from netket._src.callbacks.base import AbstractCallback, StopRun
+from netket._src.callbacks.base import (
+    AbstractCallback,
+    StopRun,
+    STOPPING_CALLBACK_ORDER,
+)
 
 
 class Timeout(AbstractCallback, mutable=True):
@@ -47,11 +51,40 @@ class Timeout(AbstractCallback, mutable=True):
         self.timeout = timeout
         self._init_time = None
 
+    @property
+    def callback_order(self) -> int:
+        # Run last, so raising StopRun never skips a later callback's collective.
+        return STOPPING_CALLBACK_ORDER
+
     def on_run_start(self, step, driver):
         self._init_time = time.time()
 
     def on_step_end(self, step, log_data, driver):
-        if time.time() - self._init_time >= self.timeout:
+        elapsed = time.time() - self._init_time
+        # Wall clocks and per-process start instants differ across nodes, so the
+        # raw per-rank decision can disagree and (without this) stop the run on
+        # different steps. Take process 0's decision on every rank, so the whole
+        # run times out in lock-step on a single authoritative clock.
+        if _timed_out_on_master(elapsed >= self.timeout):
             raise StopRun(
                 f"Timeout: training stopped after {self.timeout:.1f} seconds."
             )
+
+
+def _timed_out_on_master(local_timed_out: bool) -> bool:
+    """Return process 0's timeout decision on every process.
+
+    A no-op returning ``local_timed_out`` unchanged in single-process runs.
+    """
+    import jax
+
+    if jax.process_count() == 1:
+        return local_timed_out
+
+    import numpy as np
+    from jax.experimental import multihost_utils
+
+    out = multihost_utils.broadcast_one_to_all(
+        np.asarray(int(local_timed_out), dtype=np.int32)
+    )
+    return bool(int(out))

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

@@ -479,7 +479,15 @@ class AbstractDriver(struct.Pytree, mutable=True):
                 callbacks.on_run_end(self.step_count, self)
             except StopRun as error:
                 callbacks.on_run_end(self.step_count, self)
-                print("Stopping early because of : ", error)
+                # 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:
+                    print(
+                        f"[process {jax.process_index()}/{jax.process_count()}] "
+                        f"Stopping early because of : {error}"
+                    )
+                else:
+                    print("Stopping early because of : ", error)
 
             except KeyboardInterrupt as error:
                 callbacks.on_run_error(self.step_count, error, self)

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

@@ -357,14 +357,16 @@ class VMC_SR(AbstractOptimizationDriver):
             use_ntk = variational_state.n_parameters > variational_state.n_samples
             if jax.process_index() == 0:
                 print(
-                    "Automatic SR implementation choice: ", "NTK" if use_ntk else "QGT"
+                    f"Automatic SR implementation choice: {'NTK' if use_ntk else 'QGT'}"
+                    f" (n_parameters={variational_state.n_parameters}, n_samples={variational_state.n_samples})"
                 )
 
         if on_the_fly is None:
-            if use_ntk:
-                on_the_fly = True
-            else:
-                on_the_fly = False
+            on_the_fly = use_ntk
+            if jax.process_index() == 0:
+                print(
+                    f"Automatic SR on_the_fly choice: {'True' if on_the_fly else 'False'}"
+                )
         elif on_the_fly and not use_ntk:
             raise ValueError(
                 """

{netket-3.22.3 → netket-3.22.4}/netket/_src/logging/mlflow_log.py

@@ -15,6 +15,7 @@
 import copy
 import os
 import tempfile
+from collections.abc import Collection
 from numbers import Number
 from typing import Any, TYPE_CHECKING
 
@@ -30,6 +31,12 @@ if TYPE_CHECKING:
     from netket.vqs import VariationalState
 
 
+def _expand_node(root, tree, **kwargs):
+    if hasattr(tree, "ndim") and tree.ndim == 1:
+        return enumerate(tree)
+    return None
+
+
 def _visit_leaf(root, tree, *, data):
     if isinstance(tree, complex):
         data.append((root + "/re", tree.real))
@@ -84,6 +91,10 @@ class MLFlowLog(AbstractCallback):
             steps.  Only relevant when ``save_params=True``.
         metadata: Optional flat dict of key/value pairs logged as MLflow
             params at the start of the run.
+        ignore: Optional collection of top-level log-dict keys to skip
+            entirely.  Useful for suppressing entries that cannot be
+            serialised as MLflow metrics (e.g. large arrays or custom
+            objects).
 
     .. tip::
         Use ``metadata`` to attach a flat dict of hyper-parameters (learning rate,
@@ -143,6 +154,7 @@ class MLFlowLog(AbstractCallback):
     _tags: Any = struct.field(pytree_node=False, serialize=False)
     _save_params: Any = struct.field(pytree_node=False, serialize=False)
     _save_params_every: Any = struct.field(pytree_node=False, serialize=False)
+    _ignore: Any = struct.field(pytree_node=False, serialize=False)
     _run: Any = struct.field(pytree_node=False, serialize=False)
     _old_step: Any = struct.field(pytree_node=False, serialize=False)
     _steps_since_last_params: Any = struct.field(pytree_node=False, serialize=False)
@@ -156,6 +168,7 @@ class MLFlowLog(AbstractCallback):
         save_params: bool = True,
         save_params_every: int = 50,
         metadata: dict | None = None,
+        ignore: Collection[str] | None = None,
     ):
         super().__init__()
         self._metadata = metadata or {}
@@ -170,6 +183,7 @@ class MLFlowLog(AbstractCallback):
         self._tags = tags
         self._save_params = save_params
         self._save_params_every = save_params_every
+        self._ignore = frozenset(ignore) if ignore is not None else frozenset()
 
         if save_params and save_params_every < 1:
             raise ValueError("`save_params_every` must be >= 1")
@@ -201,7 +215,11 @@ class MLFlowLog(AbstractCallback):
             self._init_mlflow()
 
         data: list[tuple[str, Any]] = []
-        walk_tree_with_path(item, "", visit_leaf=_visit_leaf, data=data)
+        if self._ignore:
+            item = {k: v for k, v in item.items() if k not in self._ignore}
+        walk_tree_with_path(
+            item, "", visit_leaf=_visit_leaf, expand_node=_expand_node, data=data
+        )
 
         metrics = {}
         for key, val in data:

{netket-3.22.3 → netket-3.22.4}/netket/_src/logging/tensorboard.py

@@ -34,8 +34,20 @@ def _collect_tensorboard_leaf(root, tree, data):
 
 
 def _write_tensorboard_leaf(root, tree, *, writer, step):
+    key = root.removeprefix("/")
     if isinstance(tree, Number):
-        writer.add_scalar(root.removeprefix("/"), tree, step)
+        writer.add_scalar(key, tree, step)
+    elif hasattr(tree, "ndim"):
+        # numpy/jax scalar arrays (e.g. jaxlib ArrayImpl) are not `Number`
+        # instances and would otherwise be silently dropped. Only scalars
+        # (0-dim or single element) map to a tensorboard scalar.
+        if tree.ndim > 0 and tree.size != 1:
+            return
+        if jax.numpy.iscomplexobj(tree):
+            writer.add_scalar(key + "/re", float(tree.real), step)
+            writer.add_scalar(key + "/im", float(tree.imag), step)
+        else:
+            writer.add_scalar(key, float(tree), step)
 
 
 def _expand_tensorboard_node(_root, tree, **_kwargs):

{netket-3.22.3 → netket-3.22.4}/netket/_src/operator/pyscf_api.py

@@ -108,7 +108,7 @@ def from_pyscf_molecule(
         >>>
         >>> # compute the boys orbitals
         >>> mf = scf.RHF(mol).run(verbose=0)  # doctest:+ELLIPSIS
-        >>> mo_coeff = lo.Boys(mol).kernel(mf.mo_coeff)
+        >>> mo_coeff = lo.Boys(mol, mf.mo_coeff).kernel()
         >>> # use the boys orbitals to construct the netket hamiltonian
         >>> ha = nkx.operator.from_pyscf_molecule(mol, mo_coeff=mo_coeff)
 

{netket-3.22.3 → netket-3.22.4}/netket/_src/operator/pyscf_utils.py

@@ -225,7 +225,7 @@ def TV_from_pyscf_molecule(
         >>>
         >>> # compute the boys orbitals
         >>> mf = scf.RHF(mol).run(verbose=0)  # doctest:+ELLIPSIS
-        >>> mo_coeff = lo.Boys(mol).kernel(mf.mo_coeff)
+        >>> mo_coeff = lo.Boys(mol, mf.mo_coeff).kernel()
         >>> ha = nkx.operator.pyscf.TV_from_pyscf_molecule(mol, mo_coeff)
 
     Args:

{netket-3.22.3 → netket-3.22.4}/netket/_src/symmetry/representation.py

@@ -196,7 +196,7 @@ class Representation:
         group_order = len(self.group.elems)
 
         irrep_count = character_table @ self.character / group_order
-        irrep_dims = np.round(irrep_count * character_table[:, 0]).astype(int)
+        irrep_dims = np.round(np.real(irrep_count * character_table[:, 0])).astype(int)
         return irrep_dims
 
     def symmetry_adapted_basis(self) -> tuple[np.ndarray, np.ndarray]: