NetKet 3.12.1__tar.gz → 3.12.2__tar.gz

{netket-3.12.1 → netket-3.12.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.12.1
+Version: 3.12.2
 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/#
@@ -25,7 +25,7 @@ Requires-Dist: igraph<0.12.0,>=0.10.0
 Requires-Dist: jax<0.5,>=0.4.16
 Requires-Dist: jaxlib<0.5,>=0.4.16
 Requires-Dist: numba4jax<0.1,>=0.0.10.post1
-Requires-Dist: numba<0.60,>=0.52
+Requires-Dist: numba<0.61,>=0.52
 Requires-Dist: numpy~=1.20
 Requires-Dist: optax<0.3,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
@@ -63,7 +63,7 @@ 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.5,>=0.3.9; extra == 'mpi'
+Requires-Dist: mpi4jax<0.6,>=0.3.9; extra == 'mpi'
 Requires-Dist: mpi4py<4,>=3.0.1; extra == 'mpi'
 Provides-Extra: pyscf
 Requires-Dist: pyscf>=2.0; extra == 'pyscf'

{netket-3.12.1 → netket-3.12.2}/netket/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.12.1'
-__version_tuple__ = version_tuple = (3, 12, 1)
+__version__ = version = '3.12.2'
+__version_tuple__ = version_tuple = (3, 12, 2)

{netket-3.12.1 → netket-3.12.2}/netket/experimental/driver/vmc_srt.py

@@ -90,6 +90,11 @@ def SRt(
             matrix_side
         )  # * shift diagonal regularization
         aus_vector = solver_fn(matrix, dv)
+        # some solvers return a tuple, some others do not.
+        # We check and try to support both
+        if isinstance(aus_vector, tuple):
+            aus_vector, _ = aus_vector
+
         aus_vector = aus_vector.reshape(mpi.n_nodes, -1)
         aus_vector, token = mpi.mpi_scatter_jax(aus_vector, token=token)
     else:

{netket-3.12.1 → netket-3.12.2}/netket/jax/_utils_random.py

@@ -17,7 +17,7 @@ from typing import Optional
 
 import jax
 
-from netket.utils import random_seed, mpi
+from netket.utils import random_seed, mpi, config
 from netket.utils.mpi import MPI_jax_comm
 from netket.utils.types import PRNGKeyT, SeedT
 
@@ -30,16 +30,29 @@ def PRNGKey(
     The same seed will be distributed to all processes.
     """
     if seed is None:
-        key = jax.random.PRNGKey(random_seed())
-    elif isinstance(seed, int):
+        seed = random_seed()
+
+    if isinstance(seed, int):
+        # We can't sync the PRNGKey, so we can only sinc integer seeds
+        # see https://github.com/google/jax/pull/16511
+        if config.netket_experimental_sharding and jax.process_count() > 1:
+            # TODO: use stable jax function
+            from jax.experimental import multihost_utils
+
+            seed = int(
+                multihost_utils.broadcast_one_to_all(
+                    seed, is_source=jax.process_index() == root
+                ).item()
+            )
+
         key = jax.random.PRNGKey(seed)
     else:
         key = seed
 
-    key = jax.tree_util.tree_map(
-        lambda k: mpi.mpi_bcast_jax(k, root=root, comm=comm)[0], key
-    )
-
+    if not config.netket_experimental_sharding:
+        key = jax.tree_util.tree_map(
+            lambda k: mpi.mpi_bcast_jax(k, root=root, comm=comm)[0], key
+        )
     return key
 
 

{netket-3.12.1 → netket-3.12.2}/netket/jax/sharding.py

@@ -283,13 +283,17 @@ def gather(x):
     elif isinstance(x.sharding, jax.sharding.GSPMDSharding):
         # x.sharding.device_set has arbitrary order
         # Hardcode all devices until I figure out a way to deduce the order from x
-        out_shardings = PositionalSharding(jax.devices()).replicate()
+        out_shardings = (
+            PositionalSharding(jax.devices()).replicate().reshape((1,) * x.ndim)
+        )
+        # TODO support gspmdsharding in numba wrapper and use this
         # out_shardings = x.sharding.get_replicated(jax.devices())
     elif isinstance(x.sharding, PositionalSharding):
         out_shardings = x.sharding.replicate()
     else:
         raise NotImplementedError(
-            f"Gather is only compatible with PositionalSharding, but array has {x.sharding} Please open a feature request."
+            "Gather is only compatible with PositionalSharding and GSPMDSharding,"
+            f" but array has {x.sharding}. Please open a feature request."
         )
     return jax.jit(_identity, out_shardings=out_shardings)(x)
     # TODO support gspmdsharding in numba wrapper and use this

{netket-3.12.1 → netket-3.12.2}/netket/logging/json_log.py

@@ -208,7 +208,7 @@ class JsonLog(RuntimeLog):
             if self._steps_notflushed_write > 0:
                 self.flush()
         if hasattr(self, "_steps_notflushed_pars"):
-            if self._steps_notflushed_pars > 0:
+            if self._steps_notflushed_pars > 0 and self._save_params:
                 self.flush()
 
     def __repr__(self):

{netket-3.12.1 → netket-3.12.2}/netket/optimizer/preconditioner.py

@@ -15,11 +15,10 @@
 from typing import Callable, Optional, Any
 
 import abc
-from dataclasses import dataclass
 from textwrap import dedent
 
 from netket.utils.types import PyTree, Scalar
-from netket.utils import warn_deprecation, timing
+from netket.utils import warn_deprecation, timing, struct
 from netket.vqs import VariationalState
 
 from .linear_operator import LinearOperator, SolverT
@@ -43,8 +42,7 @@ def identity_preconditioner(
     return gradient
 
 
-@dataclass
-class AbstractLinearPreconditioner:
+class AbstractLinearPreconditioner(struct.Pytree, mutable=True):
     """Base class for a Linear Preconditioner solving a system :math:`Sx = F`.
 
     A LinearPreconditioner modifies the gradient :math:`F` in such a way that the new
@@ -66,7 +64,7 @@ class AbstractLinearPreconditioner:
 
     """
 
-    solver: SolverT
+    solver: SolverT = struct.field(serialize=False)
     """Function used to solve the linear system."""
 
     solver_restart: bool = False
@@ -76,13 +74,23 @@ class AbstractLinearPreconditioner:
     x0: Optional[PyTree] = None
     """Solution of the last linear system solved."""
 
-    info: Any = None
+    info: Any = struct.field(serialize=False, default=None)
     """Additional information returned by the solver when solving the last linear system."""
 
-    _lhs: LinearOperator = None
+    _lhs: LinearOperator = struct.field(serialize=False, default=None)
     """LHS of the last linear system solved."""
 
     def __init__(self, solver, *, solver_restart=False):
+        """
+        Constructs the structure holding the parameters for using the
+        linear preconditioner.
+
+        Args:
+            solver: A callable that solves a linear system of equations.
+            solver_restart: If False uses the last solution of the linear
+                system as a starting point for the solution of the next
+                (default=False).
+        """
         self.solver = solver
         self.solver_restart = solver_restart
 
@@ -119,8 +127,7 @@ class AbstractLinearPreconditioner:
 
 
 # This exists for backward compatibility
-@dataclass
-class LinearPreconditioner(AbstractLinearPreconditioner):
+class LinearPreconditioner(AbstractLinearPreconditioner, mutable=True):
     lhs_constructor: LHSConstructorT
     """Constructor of the LHS of the linear system starting from the variational state."""
 

{netket-3.12.1 → netket-3.12.2}/netket/optimizer/sr.py

@@ -16,17 +16,15 @@ from typing import Callable, Optional
 
 import jax
 
-from dataclasses import dataclass
-
 from netket.vqs import VariationalState
 from netket.utils.types import Scalar, ScalarOrSchedule
+from netket.utils import struct
 
 from .qgt import QGTAuto
 from .preconditioner import AbstractLinearPreconditioner
 
 
-@dataclass
-class SR(AbstractLinearPreconditioner):
+class SR(AbstractLinearPreconditioner, mutable=True):
     r"""
     Stochastic Reconfiguration or Natural Gradient preconditioner for the gradient.
 
@@ -63,18 +61,18 @@ class SR(AbstractLinearPreconditioner):
 
     """
 
-    diag_shift: ScalarOrSchedule = 0.01
+    diag_shift: ScalarOrSchedule = struct.field(serialize=False, default=0.01)
     """Diagonal shift added to the S matrix. Can be a Scalar value, an
        `optax <https://optax.readthedocs.io>`_ schedule or a Callable function."""
 
-    diag_scale: Optional[ScalarOrSchedule] = None
+    diag_scale: Optional[ScalarOrSchedule] = struct.field(serialize=False, default=None)
     """Diagonal shift added to the S matrix. Can be a Scalar value, an
        `optax <https://optax.readthedocs.io>`_ schedule or a Callable function."""
 
-    qgt_constructor: Callable = None
+    qgt_constructor: Callable = struct.static_field(default=None)
     """The Quantum Geometric Tensor type or a constructor."""
 
-    qgt_kwargs: dict = None
+    qgt_kwargs: dict = struct.field(serialize=False, default=None)
     """The keyword arguments to be passed to the Geometric Tensor constructor."""
 
     def __init__(

{netket-3.12.1 → netket-3.12.2}/netket/utils/mpi/gpu_autorank_util.py

@@ -42,9 +42,9 @@ def autoset_default_gpu(COMM: Optional["mpi4py.MPI.Intracomm"], verbose: bool =
     code does not run.
     """
     devices = jax.devices()
-    if len(devices) > 1 and any(d.device_kind == "gpu" for d in devices):
+    if len(devices) > 1 and any(d.platform == "gpu" for d in devices):
         local_rank = get_local_rank(COMM)
-        jax.config.set("jax_default_device", devices[local_rank])
+        jax.config.update("jax_default_device", devices[local_rank])
         logger = logging.getLogger()
         logger.info(
             f"Determined that rank {COMM.Get_rank()} will be using GPU[{local_rank}/{len(devices)}]"

{netket-3.12.1 → netket-3.12.2}/netket/utils/mpi/primitives.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import numpy as np
+import jax
 import jax.numpy as jnp
 
 from .mpi import n_nodes, MPI, MPI_py_comm, MPI_jax_comm
@@ -271,11 +272,16 @@ def mpi_bcast_jax(x, *, token=None, root, comm=MPI_jax_comm):
         return mpi4jax.bcast(x, token=token, root=root, comm=comm)
 
 
-def mpi_allgather(x, *, token=None, comm=MPI_py_comm):
+def mpi_allgather(x, *, comm=MPI_py_comm):
     if n_nodes == 1:
-        return x, token
+        return x
     else:
-        return comm.allgather(x)
+        if isinstance(x, (np.ndarray, jax.Array)):
+            out = np.empty((n_nodes,) + x.shape, dtype=x.dtype)
+            comm.Allgather(np.asarray(x), out)
+            return out
+        else:
+            return comm.allgather(x)
 
 
 def mpi_gather_jax(x, *, token=None, root: int = 0, comm=MPI_jax_comm):

{netket-3.12.1 → netket-3.12.2}/netket/utils/seed.py

@@ -19,4 +19,4 @@ def random_seed():
     """
     generates a random seed
     """
-    return np.random.default_rng().integers(0, 1 << 32)
+    return int(np.random.default_rng().integers(0, 1 << 32))

{netket-3.12.1 → netket-3.12.2}/netket/utils/struct/fields.py

@@ -51,17 +51,19 @@ jax.tree_util.register_pytree_node(
 )
 
 
-def field(pytree_node=True, serialize=True, cache=False, **kwargs):
+def field(pytree_node=True, serialize=None, cache=False, **kwargs):
     """Mark a field of a dataclass or PyTree to be:
 
     Args:
         pytree_node: a leaf node in the pytree representation of this dataclass.
             If False this must be hashable
         serialize: If True the node is included in the serialization.
-            In general you should not specify this.
+            In general you should not specify this. (Defaults to value of pytree_node)
         cache: If True this node is a cache and will be reset every time
             fields are modified.
     """
+    if serialize is None:
+        serialize = pytree_node
     return dataclasses.field(
         metadata={"pytree_node": pytree_node, "serialize": serialize, "cache": cache},
         **kwargs,

{netket-3.12.1 → netket-3.12.2}/netket/utils/struct/pytree.py

@@ -1,6 +1,7 @@
 import dataclasses
 import inspect
 import typing as tp
+from typing import Any, Callable
 from abc import ABCMeta
 from copy import copy
 from functools import partial
@@ -20,6 +21,10 @@ store the topmost non-dataclass class in a mro.
 """
 
 
+def identity(x):
+    return x
+
+
 class PytreeMeta(ABCMeta):
     """
     Metaclass for PyTrees, takes care of initializing and turning
@@ -28,13 +33,25 @@ class PytreeMeta(ABCMeta):
 
     def __call__(cls: type[P], *args: tp.Any, **kwargs: tp.Any) -> P:
         obj: P = cls.__new__(cls, *args, **kwargs)
+
+        # set things before init
+        vars_dict = vars(obj)
+
+        # Set 'default' values
+        for field in obj._pytree__cachedprop_fields:
+            vars_dict[field] = Uninitialized
+
+        for field, default_fn in cls._pytree__default_setters.items():
+            if field not in vars_dict:
+                vars_dict[field] = default_fn()
+
         obj.__dict__["_pytree__initializing"] = True
         try:
             obj.__init__(*args, **kwargs)
         finally:
             del obj.__dict__["_pytree__initializing"]
 
-        vars_dict = vars(obj)
+        # set things after init
         if obj._pytree__class_dynamic_nodes:
             vars_dict["_pytree__node_fields"] = tuple(
                 sorted(
@@ -46,9 +63,6 @@ class PytreeMeta(ABCMeta):
         else:
             vars_dict["_pytree__node_fields"] = cls._pytree__data_fields
 
-        for field in obj._pytree__cachedprop_fields:
-            vars_dict[field] = Uninitialized
-
         return obj
 
 
@@ -126,12 +140,15 @@ class Pytree(metaclass=PytreeMeta):
         class_vars = vars(cls)
         setter_descriptors = set()
         static_fields = _inherited_static_fields(cls)
+        noserialize_fields = _inherited_noserialize_fields(cls)
 
-        # add special static fields
+        # add special static fields that are in the instance, not in the class
         static_fields.add("_pytree__node_fields")
+        noserialize_fields.add("_pytree__node_fields")
 
         # new
         data_fields = _inherited_data_fields(cls)
+        default_setters = _inherited_defaults_setters(cls)  # _pytree__default_setters
         cached_prop_fields = set()
 
         for field, value in class_vars.items():
@@ -146,6 +163,18 @@ class Pytree(metaclass=PytreeMeta):
             ):
                 data_fields.add(field)
 
+            if isinstance(value, dataclasses.Field) and not value.metadata.get(
+                "serialize", True
+            ):
+                noserialize_fields.add(field)
+
+            # default setters
+            if isinstance(value, dataclasses.Field):
+                if value.default is not dataclasses.MISSING:
+                    default_setters[field] = partial(identity, value.default)
+                elif value.default_factory is not dataclasses.MISSING:
+                    default_setters[field] = value.default_factory
+
             # add setter descriptors
             if hasattr(value, "__set__"):
                 setter_descriptors.add(field)
@@ -187,14 +216,18 @@ class Pytree(metaclass=PytreeMeta):
         data_fields = tuple(sorted(data_fields))
         cached_prop_fields = tuple(sorted(cached_prop_fields))
         cached_prop_fields = tuple(_cache_name(f) for f in cached_prop_fields)
+        default_setters = MappingProxyType(default_setters)
 
         static_fields = tuple(sorted(static_fields))
+        noserialize_fields = tuple(sorted(noserialize_fields))
 
         # init class variables
         cls._pytree__initializing = False
         cls._pytree__class_is_mutable = mutable
         cls._pytree__static_fields = static_fields
+        cls._pytree__noserialize_fields = noserialize_fields
         cls._pytree__setter_descriptors = frozenset(setter_descriptors)
+        cls._pytree__default_setters = default_setters
 
         # new
         cls._pytree__class_dynamic_nodes = dynamic_nodes
@@ -234,8 +267,8 @@ class Pytree(metaclass=PytreeMeta):
 
         serialization.register_serialization_state(
             cls,
-            partial(cls._to_flax_state_dict, cls._pytree__static_fields),
-            partial(cls._from_flax_state_dict, cls._pytree__static_fields),
+            partial(cls._to_flax_state_dict, cls._pytree__noserialize_fields),
+            partial(cls._from_flax_state_dict, cls._pytree__noserialize_fields),
         )
 
     def __pre_init__(self, *args, **kwargs):
@@ -320,21 +353,21 @@ class Pytree(metaclass=PytreeMeta):
 
     @classmethod
     def _to_flax_state_dict(
-        cls, static_field_names: tuple[str, ...], pytree: "Pytree"
+        cls, noserialize_field_names: tuple[str, ...], pytree: "Pytree"
     ) -> dict[str, tp.Any]:
         from flax import serialization
 
         state_dict = {
             name: serialization.to_state_dict(getattr(pytree, name))
             for name in pytree.__dict__
-            if name not in static_field_names
+            if name not in noserialize_field_names
         }
         return state_dict
 
     @classmethod
     def _from_flax_state_dict(
         cls,
-        static_field_names: tuple[str, ...],
+        noserialize_field_names: tuple[str, ...],
         pytree: P,
         state: dict[str, tp.Any],
     ) -> P:
@@ -344,7 +377,7 @@ class Pytree(metaclass=PytreeMeta):
         state = state.copy()  # copy the state so we can pop the restored fields.
         updates = {}
         for name in pytree.__dict__:
-            if name in static_field_names:
+            if name in noserialize_field_names:
                 continue
             if name not in state:
                 raise ValueError(
@@ -435,6 +468,23 @@ def _inherited_static_fields(cls: type) -> set[str]:
     return static_fields
 
 
+def _inherited_noserialize_fields(cls: type) -> set[str]:
+    """
+    Returns the set of noserialize fields of base classes
+    of the input class
+    """
+    noserialize_fields = set()
+    for parent_class in cls.mro():
+        if parent_class is not cls and parent_class is not Pytree:
+            if issubclass(parent_class, Pytree):
+                noserialize_fields.update(parent_class._pytree__noserialize_fields)
+            elif dataclasses.is_dataclass(parent_class):
+                for field in dataclasses.fields(parent_class):
+                    if not field.metadata.get("serialize", True):
+                        noserialize_fields.add(field.name)
+    return noserialize_fields
+
+
 def _inherited_data_fields(cls: type) -> set[str]:
     """
     Returns the set of data fields of base classes
@@ -470,3 +520,23 @@ def _inherited_cachedproperty_fields(cls: type) -> set[str]:
                     if isinstance(field, CachedProperty):
                         cachedproperty_fields.add(field.name)
     return cachedproperty_fields
+
+
+def _inherited_defaults_setters(cls: type) -> set[Callable[[], Any]]:
+    """
+    Returns the set of default setters of base classes
+    of the input class.
+    """
+    default_setters = {}
+    for parent_class in cls.mro():
+        if parent_class is not cls and parent_class is not Pytree:
+            if issubclass(parent_class, Pytree):
+                default_setters.update(parent_class._pytree__default_setters)
+            elif dataclasses.is_dataclass(parent_class):
+                for field in dataclasses.fields(parent_class):
+                    if field.default is not dataclasses.MISSING:
+                        default_setters[field.name] = lambda: field.default
+                    elif field.default_factory is not dataclasses.MISSING:
+                        default_setters[field.name] = field.default_factory
+
+    return default_setters

{netket-3.12.1 → netket-3.12.2}/pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     "scipy>=1.5.3, <2",
     "tqdm>=4.60, <5",
     "rich>=10.0",
-    "numba>=0.52, <0.60",
+    "numba>=0.52, <0.61",
     "igraph>=0.10.0, <0.12.0",
     "jax>=0.4.16, <0.5",
     "jaxlib>=0.4.16, <0.5",
@@ -56,7 +56,7 @@ dependencies = [
 [project.optional-dependencies]
 mpi = [
     "mpi4py>=3.0.1, <4",
-    "mpi4jax>=0.3.9, <0.5"
+    "mpi4jax>=0.3.9, <0.6"
     ]
 pyscf = [
     "sparse>=0.12.0",