NetKet 3.20.0__tar.gz → 3.20.2__tar.gz

{netket-3.20.0 → netket-3.20.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.20.0
+Version: 3.20.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/#
@@ -46,7 +46,7 @@ Requires-Dist: pytest-json-report>=1.3; extra == 'dev'
 Requires-Dist: pytest-xdist[psutil]>=2; extra == 'dev'
 Requires-Dist: pytest>=6; extra == 'dev'
 Requires-Dist: qutip>=5; extra == 'dev'
-Requires-Dist: ruff==0.12.2; extra == 'dev'
+Requires-Dist: ruff==0.14.1; extra == 'dev'
 Requires-Dist: wheel>=0.38.1; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: ipywidgets>=8.1.7; extra == 'docs'

{netket-3.20.0 → netket-3.20.2}/netket/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.20.0'
-__version_tuple__ = version_tuple = (3, 20, 0)
+__version__ = version = '3.20.2'
+__version_tuple__ = version_tuple = (3, 20, 2)
 
 __commit_id__ = commit_id = None

{netket-3.20.0 → netket-3.20.2}/netket/jax/__init__.py

@@ -23,7 +23,7 @@ from ._utils_dtype import (
 )
 
 from ._utils_tree import (
-    tree_ravel,
+    tree_ravel as _deprecated_tree_ravel,
     tree_size,
     eval_shape,
     tree_leaf_isreal,
@@ -69,3 +69,18 @@ from . import sharding
 from netket.utils import _hide_submodules
 
 _hide_submodules(__name__, ignore="sharding")
+
+# Deprecation machinery
+_deprecations = {
+    # January 2025, NetKet 3.16
+    "tree_ravel": (
+        "netket.jax.tree_ravel is deprecated: use jax.flatten_util.ravel_pytree directly",
+        _deprecated_tree_ravel,
+    ),
+}
+
+from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
+
+__getattr__ = _deprecation_getattr(__name__, _deprecations)
+
+del _deprecation_getattr

{netket-3.20.0 → netket-3.20.2}/netket/jax/_jacobian/jacobian_dense.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from jax import numpy as jnp
+from jax.flatten_util import ravel_pytree
 
 from netket.utils.types import Array, PyTree
 import netket.jax as nkjax
@@ -25,9 +26,9 @@ from .jacobian_pytree import (
 
 def ravel(x: PyTree) -> Array:
     """
-    shorthand for tree_ravel
+    shorthand for ravel_pytree
     """
-    dense, _ = nkjax.tree_ravel(x)
+    dense, _ = ravel_pytree(x)
     return dense
 
 

{netket-3.20.0 → netket-3.20.2}/netket/jax/_jacobian/logic.py

@@ -121,8 +121,8 @@ def jacobian(
             dimension matches the number of samples (the first axis of **samples**),
             while the second dimension will match the total number of parameters.
             This raveling is equivalent to :func:`jax.vmap` of
-            :func:`netket.jax.tree_ravel`,
-            :code:`jax.vmap(nk.jax.tree_ravel, nk.jax.jacobian(...))`.
+            :func:`jax.flatten_util.ravel_pytree`,
+            :code:`jax.vmap(ravel_pytree, nk.jax.jacobian(...))`.
             If using **complex** mode with real parameters the returned tensor
             has 3 dimensions, where the first and last match the other modes while
             the middle one has size 2, and encodes the gradient of the real and

{netket-3.20.0 → netket-3.20.2}/netket/jax/_utils_tree.py

@@ -17,15 +17,13 @@ from collections.abc import Callable
 
 
 import jax
-import netket.jax as nkjax
 from jax import numpy as jnp
 from jax.tree_util import (
     register_pytree_node,
-    tree_flatten,
-    tree_unflatten,
     tree_map,
     tree_leaves,
 )
+from jax.flatten_util import ravel_pytree
 
 from netket.utils.types import PyTree, Scalar
 from netket.utils.numbers import is_scalar
@@ -43,14 +41,7 @@ def tree_ravel(pytree: PyTree) -> tuple[jnp.ndarray, Callable]:
       unflattening a 1D vector of the same length back to a pytree of of the same
       structure as the input ``pytree``.
     """
-    leaves, treedef = tree_flatten(pytree)
-    flat, unravel_list = nkjax.vjp(_ravel_list, *leaves)
-    unravel_pytree = lambda flat: tree_unflatten(treedef, unravel_list(flat))
-    return flat, unravel_pytree
-
-
-def _ravel_list(*lst):
-    return jnp.concatenate([jnp.ravel(elt) for elt in lst]) if lst else jnp.array([])
+    return ravel_pytree(pytree)
 
 
 def eval_shape(fun, *args, has_aux=False, **kwargs):
@@ -131,7 +122,7 @@ def tree_norm(a: PyTree, ord: int = 2) -> Scalar:
     r"""
     Compute the norm of a pytree, intended as a 1D vector of values.
 
-    Equivalent to :code:`jnp.linalg.norm(nk.jax.tree_ravel(a)[0], ord)`.
+    Equivalent to :code:`jnp.linalg.norm(jax.flatten_util.ravel_pytree(a)[0], ord)`.
 
     Args:
         a: A pytree, interpreted as a vector

{netket-3.20.0 → netket-3.20.2}/netket/optimizer/qgt/qgt_jacobian_dense.py

@@ -15,6 +15,7 @@
 
 import jax
 from jax import numpy as jnp
+from jax.flatten_util import ravel_pytree
 from flax import struct
 
 from netket.utils.types import Scalar, PyTree
@@ -286,6 +287,6 @@ def convert_tree_to_dense_format(vec, mode, *, disable=False):
         if mode != "holomorphic":
             vec, reassemble = nkjax.tree_to_real(vec)
         if not hasattr(vec, "ndim"):
-            vec, unravel = nkjax.tree_ravel(vec)
+            vec, unravel = ravel_pytree(vec)
 
     return vec, lambda x: reassemble(unravel(x))

{netket-3.20.0 → netket-3.20.2}/netket/optimizer/qgt/qgt_jacobian_pytree.py

@@ -15,6 +15,7 @@
 
 import jax
 from jax import numpy as jnp
+from jax.flatten_util import ravel_pytree
 from flax import struct
 
 from netket.utils.types import Array, PyTree, Scalar
@@ -71,7 +72,7 @@ class QGTJacobianPyTreeT(LinearOperator):
     def __matmul__(self, vec: PyTree | Array) -> PyTree | Array:
         # Turn vector RHS into PyTree
         if hasattr(vec, "ndim"):
-            _, unravel = nkjax.tree_ravel(self._params_structure)
+            _, unravel = ravel_pytree(self._params_structure)
             vec = unravel(vec)
             ravel = True
         else:
@@ -98,7 +99,7 @@ class QGTJacobianPyTreeT(LinearOperator):
 
         # Ravel PyTree back into vector as needed
         if ravel:
-            result, _ = nkjax.tree_ravel(result)
+            result, _ = ravel_pytree(result)
 
         return result
 
@@ -179,12 +180,12 @@ class QGTJacobianPyTreeT(LinearOperator):
         if self.mode == "complex" or self.mode == "imag":
             # I want to iterate across the samples and real/imaginary part
             O = jax.tree_util.tree_map(lambda x: x.reshape(-1, *x.shape[2:]), O)
-        O = jax.vmap(lambda l: nkjax.tree_ravel(l)[0])(O)
+        O = jax.vmap(lambda l: ravel_pytree(l)[0])(O)
 
         if self.scale is None:
             diag = jnp.eye(O.shape[-1])
         else:
-            scale, _ = nkjax.tree_ravel(self.scale)
+            scale, _ = ravel_pytree(self.scale)
             O = O * scale[jnp.newaxis, :]
             diag = jnp.diag(scale**2)
 

{netket-3.20.0 → netket-3.20.2}/netket/optimizer/qgt/qgt_onthefly.py

@@ -17,6 +17,7 @@ import warnings
 
 import jax
 from jax import numpy as jnp
+from jax.flatten_util import ravel_pytree
 from flax import struct
 
 import netket.jax as nkjax
@@ -248,7 +249,7 @@ def onthefly_mat_treevec(
                              """
             )
 
-        _, unravel = nkjax.tree_ravel(S._params)
+        _, unravel = ravel_pytree(S._params)
         vec = unravel(vec)
         ravel_result = True
     else:
@@ -261,7 +262,7 @@ def onthefly_mat_treevec(
     res = S._mat_vec(vec, S.diag_shift)
 
     if ravel_result:
-        res, _ = nkjax.tree_ravel(res)
+        res, _ = ravel_pytree(res)
 
     return res
 

{netket-3.20.0 → netket-3.20.2}/netket/optimizer/solver/solvers.py

@@ -16,8 +16,7 @@
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
-
-from netket.jax import tree_ravel
+from jax.flatten_util import ravel_pytree
 from netket.utils.api_utils import partial_from_kwargs
 from netket.utils.deprecation import warn_deprecation
 
@@ -100,7 +99,7 @@ def pinv_smooth(
 
     if not isinstance(A, jax.Array):
         A = A.to_dense()
-    b, unravel = tree_ravel(b)
+    b, unravel = ravel_pytree(b)
 
     Σ, U = jnp.linalg.eigh(A)
 
@@ -166,7 +165,7 @@ def pinv(A, b, *, rtol: float = 1e-12, x0=None, rcond: float = None):
 
     if not isinstance(A, jax.Array):
         A = A.to_dense()
-    b, unravel = tree_ravel(b)
+    b, unravel = ravel_pytree(b)
 
     A_inv = jnp.linalg.pinv(A, rtol=rtol, hermitian=True)
 
@@ -197,7 +196,7 @@ def svd(A, b, *, rcond=None, x0=None):
 
     if not isinstance(A, jax.Array):
         A = A.to_dense()
-    b, unravel = tree_ravel(b)
+    b, unravel = ravel_pytree(b)
 
     x, residuals, rank, s = jnp.linalg.lstsq(A, b, rcond=rcond)
 
@@ -228,7 +227,7 @@ def cholesky(A, b, *, lower=False, x0=None):
 
     if not isinstance(A, jax.Array):
         A = A.to_dense()
-    b, unravel = tree_ravel(b)
+    b, unravel = ravel_pytree(b)
 
     c, low = jsp.linalg.cho_factor(A, lower=lower)
     x = jsp.linalg.cho_solve((c, low), b)
@@ -259,7 +258,7 @@ def LU(A, b, *, trans=0, x0=None):
 
     if not isinstance(A, jax.Array):
         A = A.to_dense()
-    b, unravel = tree_ravel(b)
+    b, unravel = ravel_pytree(b)
 
     lu, piv = jsp.linalg.lu_factor(A)
     x = jsp.linalg.lu_solve((lu, piv), b, trans=0)
@@ -291,7 +290,7 @@ def solve(A, b, *, assume_a="pos", x0=None):
 
     if not isinstance(A, jax.Array):
         A = A.to_dense()
-    b, unravel = tree_ravel(b)
+    b, unravel = ravel_pytree(b)
 
     x = jsp.linalg.solve(A, b, assume_a="pos")
     return unravel(x), None

{netket-3.20.0 → netket-3.20.2}/netket/utils/model_frameworks/nnx.py

@@ -80,12 +80,18 @@ class NNXWrapper:
         return nnx_module
 
     def __getattr__(self, name):
-        if hasattr(_get_graphdef_type(self.graphdef), name):
-            return partial(self.apply, method=name)
-        raise AttributeError(
-            f"{_get_graphdef_type(self.static_module)} (wrapped into a '{type(self).__name__}')"
-            f"has no attribute '{name}'"
-        )
+        # Use try-except instead of hasattr to avoid infinite recursion
+        # when checking for attributes that don't exist
+        graphdef_type = _get_graphdef_type(object.__getattribute__(self, "graphdef"))
+        try:
+            # Check if the attribute exists on the graphdef type's class
+            getattr(graphdef_type, name)
+            return partial(object.__getattribute__(self, "apply"), method=name)
+        except AttributeError:
+            raise AttributeError(
+                f"{graphdef_type} (wrapped into a '{type(self).__name__}')"
+                f" has no attribute '{name}'"
+            )
 
     def __hash__(self):
         return hash(("NNXWrapper", self.graphdef))

{netket-3.20.0 → netket-3.20.2}/pyproject.toml

@@ -63,7 +63,7 @@ dev = [
     "coverage>=5",
     "pre-commit>=2.7",
     "black==25.1.0",
-    "ruff==0.12.2",
+    "ruff==0.14.1",
     "qutip>=5",
     "wheel>=0.38.1",
     "build>=1",