NetKet 3.17__tar.gz → 3.17.1__tar.gz

{netket-3.17 → netket-3.17.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.17
+Version: 3.17.1
 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/#

{netket-3.17 → netket-3.17.1}/netket/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.17'
-__version_tuple__ = version_tuple = (3, 17)
+__version__ = version = '3.17.1'
+__version_tuple__ = version_tuple = (3, 17, 1)

{netket-3.17 → netket-3.17.1}/netket/experimental/logging/hdf5_log.py

@@ -18,6 +18,7 @@ from flax.serialization import to_bytes
 from flax.core import pop as fpop, FrozenDict
 
 from netket.logging import AbstractLog
+from netket.jax.sharding import extract_replicated
 
 _mode_shorthands = {"write": "w", "append": "a", "fail": "x"}
 
@@ -172,22 +173,27 @@ class HDF5Log(AbstractLog):
         if self._writer is None:
             self._init_output_file()
 
+        if self._steps_notsaved_params % self._save_params_every == 0:
+            variables = variational_state.variables
+
+            # TODO: remove - FrozenDict are deprecated
+            if isinstance(variables, FrozenDict):
+                variables = variables.unfreeze()
+            variables = extract_replicated(variables)
+            self._steps_notsaved_params = 0
+        else:
+            variables = None
+
         if self._is_master_process:
             tree_log(log_data, "data", self._writer, iter=step)
 
-            if self._steps_notsaved_params % self._save_params_every == 0:
-                variables = variational_state.variables
-                # TODO: remove - FrozenDict are deprecated
-                if isinstance(variables, FrozenDict):
-                    variables = variables.unfreeze()
-
+            if variables is not None:
                 _, params = fpop(variables, "params")
                 binary_data = to_bytes(variables)
                 tree = {"model_state": binary_data, "parameters": params, "iter": step}
                 tree_log(tree, "variational_state", self._writer)
-                self._steps_notsaved_params = 0
-
             self._writer.flush()
+
         self._steps_notsaved_params += 1
 
     def flush(self, variational_state=None):

{netket-3.17 → netket-3.17.1}/netket/utils/history/history_dict.py

@@ -34,6 +34,17 @@ else:
     History = Any
 
 
+def _format_key(k):
+    if hasattr(k, "key"):
+        return str(k.key)
+    elif hasattr(k, "idx"):
+        return str(k.idx)
+    elif hasattr(k, "name"):
+        return str(k, k.name)
+    else:
+        return str(k)  # fallback
+
+
 class HistoryDict:
     """A class that behaves like a dictionary, mapping strings to History instances or nested
     HistoryDict instances.
@@ -91,7 +102,7 @@ class HistoryDict:
 
         _repr_str = f"HistoryDict with {len(dk)} elements:"
         for keys, val in dk:
-            path = "/".join([val.key for val in keys])
+            path = "/".join([_format_key(k) for k in keys])
             _repr_str += f"\n\t'{path}' -> " + str(val)
 
         return _repr_str