NetKet 3.22.2__tar.gz → 3.22.3__tar.gz

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

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

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '3.22.2'
-__version_tuple__ = version_tuple = (3, 22, 2)
+__version__ = version = '3.22.3'
+__version_tuple__ = version_tuple = (3, 22, 3)
 
 __commit_id__ = commit_id = None

{netket-3.22.2 → netket-3.22.3}/netket/utils/_tips.py

@@ -18,6 +18,8 @@ import random
 import sys
 
 TIPS = [
+    # News
+    "NetKet 3.22 has plenty of new stuff! Check it out at https://www.netket.org/posts/07-netket322/",
     # Drivers
     "Prefer the new nk.driver.VMC_SR over VMC which supports minSR and SPRING.",
     "Use driver.run(..., timeit=True) to know where your dominant cost is.",

{netket-3.22.2 → netket-3.22.3}/netket/utils/model_frameworks/base.py

@@ -31,6 +31,23 @@ class ModuleFramework(abc.ABC):
         """
         return False
 
+    @staticmethod
+    def wrapped_model_class(static_model: Any) -> type:
+        """
+        Returns the concrete user-defined model class, given the (possibly wrapped)
+        static module stored in :attr:`MCState._model`.
+
+        For flax-like frameworks the static module *is* the model, so this returns
+        its type. Frameworks that wrap the model into a static container (such as
+        ``flax.nnx`` or ``equinox``) override this to peek inside the wrapper and
+        return the original class.
+
+        This lets serialization code recover the original model class without
+        unwrapping (which would require the variables), e.g. to dispatch on whether
+        a custom serializer is registered for that class.
+        """
+        return type(static_model)
+
     @staticmethod
     @abc.abstractmethod
     def is_loaded() -> bool:

{netket-3.22.2 → netket-3.22.3}/netket/utils/model_frameworks/equinox.py

@@ -56,7 +56,14 @@ class EquinoxWrapper:
         module = self.recompose(variables)
         fun = getattr(module, method)
 
-        return fun(*args, key=rngs, **kwargs)
+        # Only forward a `key` to the model if one was actually provided.
+        # Passing `key=None` unconditionally breaks Equinox modules whose
+        # `__call__` does not declare a `key` argument (e.g. simple models
+        # that do not use randomness).
+        if rngs is not None:
+            kwargs["key"] = rngs
+
+        return fun(*args, **kwargs)
 
     def recompose(self, variables):
         import equinox as eqx
@@ -94,6 +101,11 @@ class EquinoxFramework(ModuleFramework):
         """
         return True
 
+    @staticmethod
+    def wrapped_model_class(static_model: "EquinoxWrapper") -> type:
+        # The static partition is an instance of the original equinox.Module class.
+        return type(static_model.static_module)
+
     @staticmethod
     def is_loaded() -> bool:
         return "equinox" in sys.modules

{netket-3.22.2 → netket-3.22.3}/netket/utils/model_frameworks/nnx.py

@@ -107,6 +107,11 @@ class NNXFramework(ModuleFramework):
         """
         return True
 
+    @staticmethod
+    def wrapped_model_class(static_model: "NNXWrapper") -> type:
+        # The original nnx.Module class is recorded in the graphdef.
+        return _get_graphdef_type(static_model.graphdef)
+
     @staticmethod
     def is_loaded() -> bool:
         # this should be not necessary, as netket requires and loads