NetKet 3.18__tar.gz → 3.18.1__tar.gz

{netket-3.18 → netket-3.18.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.18
+Version: 3.18.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.18 → netket-3.18.1}/netket/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.18'
-__version_tuple__ = version_tuple = (3, 18)
+__version__ = version = '3.18.1'
+__version_tuple__ = version_tuple = (3, 18, 1)

{netket-3.18 → netket-3.18.1}/netket/operator/_fermion2nd/base.py

@@ -11,11 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import warnings
+from textwrap import dedent
 
+from numbers import Number
 
 import numpy as np
 
-from numbers import Number
+import jax
 
 from netket.utils.types import DType
 from netket.operator import DiscreteOperator, Transpose
@@ -24,7 +27,8 @@ from netket.hilbert import AbstractHilbert
 from netket.utils.numbers import is_scalar, dtype as _dtype
 from netket.utils.optional_deps import import_optional_dependency
 
-from netket.hilbert import SpinOrbitalFermions
+from netket.hilbert import Fock, SpinOrbitalFermions
+
 
 from .utils import (
     _convert_terms_to_spin_blocks,
@@ -137,7 +141,39 @@ class FermionOperator2ndBase(DiscreteOperator):
 
     def _setup(self, force: bool = False):
         """Analyze the operator strings and precompute arrays for get_conn inference"""
-        raise NotImplementedError  # pragma: no cover
+        if jax.process_index() == 0:
+            if (
+                isinstance(self.hilbert, SpinOrbitalFermions)
+                and self.hilbert.n_fermions_per_spin is not None
+            ):
+                warnings.warn(
+                    dedent(
+                        """
+                        WARNING: Initializing `netket.operator.FermionOperator2nd` for a Hilbert space with a fixed number of fermions.
+                        Consider using `netket.experimental.operator.ParticleNumberAndSpinConservingFermioperator2nd` to reduce the number of connected elements and
+                        considerably reduce the computational cost.
+                        You can convert this operator by calling `netket.experimental.operator.ParticleNumberAndSpinConservingFermioperator2nd.from_fermionoperator2nd`.
+                        """
+                    ),
+                    stacklevel=2,
+                )
+            elif (
+                isinstance(self.hilbert, SpinOrbitalFermions)
+                and self.hilbert.n_fermions is not None
+            ) or (
+                isinstance(self.hilbert, Fock) and self.hilbert.n_particles is not None
+            ):
+                warnings.warn(
+                    dedent(
+                        """
+                        WARNING: Initializing `netket.operator.FermionOperator2nd` for a Hilbert space with a fixed number of fermions.
+                        Consider using `netket.experimental.operator.ParticleNumberConservingFermioperator2nd` to reduce the number of connected elements and
+                        considerably reduce the computational cost.
+                        You can convert this operator by calling `netket.experimental.operator.ParticleNumberConservingFermioperator2nd.from_fermionoperator2nd`.
+                        """
+                    ),
+                    stacklevel=2,
+                )
 
     @classmethod
     def from_openfermion(

{netket-3.18 → netket-3.18.1}/netket/operator/_fermion2nd/jax.py

@@ -558,6 +558,7 @@ class FermionOperator2ndJax(FermionOperator2ndBase, DiscreteJaxOperator):
     def _setup(self, force: bool = False):
         """Analyze the operator strings and precompute arrays for get_conn inference"""
         if force or not self._initialized:
+            super()._setup(self)
             # TODO ideally we would set dagger_dtype to the same as x
             # however, unfortunately, the dtype of the states in netket
             # is stored in the sampler and not in hilbert, so we don't know it at this stage

{netket-3.18 → netket-3.18.1}/netket/operator/_fermion2nd/numba.py

@@ -52,6 +52,7 @@ class FermionOperator2nd(FermionOperator2ndBase):
     def _setup(self, force: bool = False):
         """Analyze the operator strings and precompute arrays for get_conn inference"""
         if force or not self._initialized:
+            super()._setup(self)
             # following lists will be used to compute matrix elements
             # they are filled in _add_term
             out = _pack_internals(self._operators, self._dtype)

{netket-3.18 → netket-3.18.1}/netket/sampler/metropolis.py

@@ -314,11 +314,14 @@ class MetropolisSampler(Sampler):
         )
         n_chains_per_rank = n_chains // device_count()
 
-        if chunk_size is not None and n_chains_per_rank % chunk_size != 0:
+        if (
+            chunk_size is not None
+            and n_chains_per_rank > chunk_size
+            and n_chains_per_rank % chunk_size != 0
+        ):
             raise ValueError(
                 f"Chunk size must divide number of chains per rank, {n_chains_per_rank}"
             )
-        self.chunk_size = chunk_size
 
         super().__init__(
             hilbert=hilbert,

{netket-3.18 → netket-3.18.1}/netket/utils/model_frameworks/nnx.py

@@ -17,12 +17,24 @@ from functools import partial
 
 import sys
 
+from netket.utils.version_check import module_version
+
 from .base import ModuleFramework, framework
 
 if TYPE_CHECKING:
     from flax import nnx
 
 
+def _get_graphdef_type(graphdef):
+    # TODO: Once we require Jax 0.10.6 (probably once we require jax 0.7)
+    # drop this check
+    if module_version("flax") >= (0, 10, 6):
+        assert len(graphdef.nodes) > 0
+        return graphdef.nodes[0].type
+    else:
+        return graphdef.type
+
+
 # expose jax-stax as a flax module
 class NNXWrapper:
     name: str = "NNX"
@@ -68,14 +80,15 @@ class NNXWrapper:
         return nnx_module
 
     def __getattr__(self, name):
-        if hasattr(self.graphdef.type, name):
+        if hasattr(_get_graphdef_type(self.graphdef), name):
             return partial(self.apply, method=name)
         raise AttributeError(
             f"'{type(self).__name__}' (and the wrapped '{self.graphdef.type}') object has no attribute '{name}'"
         )
 
     def __repr__(self):
-        return f"NNXWrapper(wrapped_class={self.graphdef.type}, ...)"
+        typ = _get_graphdef_type(self.graphdef)
+        return f"NNXWrapper(wrapped_class={typ}, ...)"
 
 
 @framework

{netket-3.18 → netket-3.18.1}/netket/utils/struct/fields.py

@@ -16,7 +16,6 @@
 Common utilities to define fields and properties of netket dataclasses
 and Pytrees.
 """
-
 from functools import partial
 
 import dataclasses
@@ -55,6 +54,7 @@ jax.tree_util.register_pytree_node(
 
 def field(
     pytree_node: bool = True,
+    pytree_ignore: bool = False,
     serialize: bool | None = None,
     serialize_name: str | None = None,
     cache: bool = False,
@@ -65,7 +65,13 @@ def field(
 
     Args:
         pytree_node: a leaf node in the pytree representation of this dataclass.
-            If False this must be hashable
+            If False this must be hashable.
+        pytree_ignore: If True this field is ignored by the pytree metadata and will be
+            excluded from the pytree flattening/unflattening process. This means the field
+            will not appear in the flattened representation when calling `jax.tree_util.tree_flatten()`
+            or similar pytree operations. This is useful for caches, temporary data, or other
+            fields that should not be passed forward during pytree transformations.
+            When True, pytree_node must be False.
         serialize: If True the node is included in the serialization.
             In general you should not specify this. (Defaults to value of pytree_node).
         serialize_name: If specified, it's the name under which this attribute is serialized.
@@ -80,9 +86,14 @@ def field(
         serialize = pytree_node
     if sharded is True:
         sharded = ShardedFieldSpec()
+    if pytree_ignore and pytree_node:
+        raise ValueError(
+            "Cannot set both `pytree_node=True` and `pytree_ignore=True`. Pytreenode must be False if pytree_ignore is True."
+        )
 
     metadata = {
         "pytree_node": pytree_node,
+        "ignore": pytree_ignore,
         "serialize": serialize,
         "cache": cache,
         "sharded": sharded,
@@ -104,11 +115,12 @@ class CachedProperty:
     but must be cached.
     """
 
-    def __init__(self, method, pytree_node=False):
+    def __init__(self, method, pytree_node=False, pytree_ignore=False):
         self.name = method.__name__
         self.cache_name = _cache_name(self.name)
         self.method = method
         self.pytree_node = pytree_node
+        self.pytree_ignore = pytree_ignore
         self.type = method.__annotations__.get("return", MISSING)
         self.doc = method.__doc__
 
@@ -131,7 +143,7 @@ class CachedProperty:
         )
 
 
-def property_cached(fun=None, pytree_node=False):
+def property_cached(fun=None, pytree_node=False, pytree_ignore=False):
     """Decorator to make the method behave as a property but cache the resulting value and
     clears it upon replace.
 
@@ -139,7 +151,14 @@ def property_cached(fun=None, pytree_node=False):
         pytree_node: a leaf node in the pytree representation of this dataclass.
             If False this must be hashable
     """
+    if pytree_ignore and pytree_node:
+        raise ValueError(
+            "Cannot set both `pytree_node=True` and `pytree_ignore=True`. Pytreenode must be False if pytree_ignore is True."
+        )
+
     if fun is None:
-        return partial(property_cached, pytree_node=pytree_node)
+        return partial(
+            property_cached, pytree_node=pytree_node, pytree_ignore=pytree_ignore
+        )
 
-    return CachedProperty(fun, pytree_node=pytree_node)
+    return CachedProperty(fun, pytree_node=pytree_node, pytree_ignore=pytree_ignore)

{netket-3.18 → netket-3.18.1}/netket/utils/struct/pytree.py

@@ -137,6 +137,7 @@ class Pytree(metaclass=PytreeMeta):
     _pytree__fields: dict[str, dataclasses.Field]
     _pytree__static_fields: tuple[str, ...]
     _pytree__node_fields: tuple[str, ...]
+    _pytree__ignored_fields: tuple[str, ...]
     _pytree__setter_descriptors: frozenset[str]
 
     _pytree__cachedprop_fields: tuple[str, ...]
@@ -147,7 +148,7 @@ class Pytree(metaclass=PytreeMeta):
         # gather class info
         class_vars = vars(cls)
         setter_descriptors = set()
-        all_fields, data_fields, static_fields = _inherited_fields(cls)
+        all_fields, data_fields, static_fields, ignored_fields = _inherited_fields(cls)
 
         # new
         default_setters = _inherited_defaults_setters(cls)  # _pytree__default_setters
@@ -166,7 +167,10 @@ class Pytree(metaclass=PytreeMeta):
                 all_fields[field] = value
                 _is_pytree_node = value.metadata.get("pytree_node", True)
                 if not _is_pytree_node:
-                    static_fields.add(field)
+                    if value.metadata.get("ignore", False):
+                        ignored_fields.add(field)
+                    else:
+                        static_fields.add(field)
                 elif _is_pytree_node:
                     data_fields.add(field)
 
@@ -192,6 +196,8 @@ class Pytree(metaclass=PytreeMeta):
             # setattr(cls, _cache_name(field), Uninitialized)
             if class_vars[field].pytree_node:
                 data_fields.add(_cache_name(field))
+            elif class_vars[field].pytree_ignore:
+                ignored_fields.add(_cache_name(field))
             else:
                 static_fields.add(_cache_name(field))
         cached_prop_fields = cached_prop_fields.union(
@@ -242,13 +248,16 @@ class Pytree(metaclass=PytreeMeta):
                     pass
 
         # new
-        init_fields = tuple(sorted(data_fields.union(static_fields)))
+        init_fields = tuple(
+            sorted(data_fields.union(static_fields).union(ignored_fields))
+        )
         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))
+        ignored_fields = tuple(sorted(ignored_fields))
         noserialize_fields = tuple(sorted(noserialize_fields))
         serialize_rename_fields = tuple(
             sorted(compute_serialize_rename_fields(all_fields))
@@ -260,6 +269,7 @@ class Pytree(metaclass=PytreeMeta):
         cls._pytree__class_is_mutable = mutable
         cls._pytree__fields = all_fields
         cls._pytree__static_fields = static_fields
+        cls._pytree__ignored_fields = ignored_fields
         cls._pytree__noserialize_fields = noserialize_fields
         cls._pytree__serialize_rename_fields = serialize_rename_fields
         cls._pytree__sharded_fields = sharded_fields
@@ -357,6 +367,10 @@ class Pytree(metaclass=PytreeMeta):
             raise KeyError(
                 f"Missing static field {e} in PyTree {cls.__name__} while flattening."
             ) from e
+
+        # drop ignored keys
+        for ignored_field in pytree._pytree__ignored_fields:
+            all_vars.pop(ignored_field, None)
         if with_key_paths:
             node_values = tuple(
                 (jax.tree_util.GetAttrKey(field), all_vars.pop(field))
@@ -549,7 +563,7 @@ class Pytree(metaclass=PytreeMeta):
 
 def _inherited_fields(
     cls: type,
-) -> tuple[dict[str, dataclasses.Field], set[str], set[str]]:
+) -> tuple[dict[str, dataclasses.Field], set[str], set[str], set[str]]:
     """
     Returns the set of all fields (static and data) fields from
     base classes.
@@ -557,6 +571,7 @@ def _inherited_fields(
     fields = {}
     data_fields = set()
     static_fields = set()
+    ignored_fields = set()
 
     for parent_class in cls.mro():
         if parent_class is not cls and parent_class is not Pytree:
@@ -564,17 +579,22 @@ def _inherited_fields(
                 fields.update(parent_class._pytree__fields)
                 data_fields.update(parent_class._pytree__data_fields)
                 static_fields.update(parent_class._pytree__static_fields)
+                ignored_fields.update(parent_class._pytree__ignored_fields)
             elif dataclasses.is_dataclass(parent_class):
                 for field in dataclasses.fields(parent_class):
                     fields[field.name] = field
 
                     _is_pytree_node = field.metadata.get("pytree_node", True)
+                    _ignore_node = field.metadata.get("struct_ignore", False)
                     if _is_pytree_node:
                         data_fields.add(field.name)
                     else:
-                        static_fields.add(field.name)
+                        if _ignore_node:
+                            ignored_fields.add(field.name)
+                        else:
+                            static_fields.add(field.name)
 
-    return fields, data_fields, static_fields
+    return fields, data_fields, static_fields, ignored_fields
 
 
 def _inherited_noserialize_fields(cls: type) -> set[str]: