NetKet 3.19.0__tar.gz → 3.19.2__tar.gz

{netket-3.19.0 → netket-3.19.2}/PKG-INFO

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

{netket-3.19.0 → netket-3.19.2}/netket/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '3.19.0'
-__version_tuple__ = version_tuple = (3, 19, 0)
+__version__ = version = '3.19.2'
+__version_tuple__ = version_tuple = (3, 19, 2)
+
+__commit_id__ = commit_id = None

{netket-3.19.0 → netket-3.19.2}/netket/graph/graph.py

@@ -181,7 +181,7 @@ class Graph(AbstractGraph):
 
     def _compute_automorphisms(self):
         """
-        Compute the graph autmorphisms of this graph.
+        Compute the graph automorphisms of this graph.
         """
         colors = self.edge_colors
         result = self._igraph.get_isomorphisms_vf2(
@@ -190,7 +190,10 @@ class Graph(AbstractGraph):
 
         # sort them s.t. the identity comes first
         result = np.unique(result, axis=0).tolist()
-        result = PermutationGroup([Permutation(i) for i in result], self.n_nodes)
+        result = PermutationGroup(
+            [Permutation(inverse_permutation_array=perm) for perm in result],
+            self.n_nodes,
+        )
         return result
 
     # TODO turn into a struct.property_cached?

{netket-3.19.0 → netket-3.19.2}/netket/hilbert/index/constrained_sum.py

@@ -78,7 +78,8 @@ class SumConstrainedHilbertIndex(HilbertIndex):
     @jax.jit
     def _compute_all_states(self):
         if self.n_particles == 0:
-            return jnp.zeros((1, self.size), dtype=self.range.dtype)
+            all_states_fock = jnp.zeros((1, self.size), dtype=self.range.dtype)
+            return self.range.numbers_to_states(all_states_fock, dtype=self.range.dtype)
         with jax.ensure_compile_time_eval():
             c = jnp.repeat(
                 jnp.eye(self.size, dtype=self.range.dtype),

{netket-3.19.0 → netket-3.19.2}/netket/optimizer/qgt/qgt_onthefly_logic.py

@@ -90,77 +90,93 @@ def mat_vec_factory(forward_fn, params, model_state, samples, pdf=None):
 # Methods below are needed for the chunked version of QGTOnTheFly
 
 
-@partial(sharding_decorator, sharded_args_tree=(False, False, True, False, False))
-def _O_jvp(forward_fn, params, samples, v, chunk_size):
-    @partial(scanmap, scan_fun=scan_append, argnums=2)
-    def __O_jvp(forward_fn, params, samples, v):
+@partial(
+    sharding_decorator, sharded_args_tree=(False, False, False, True, False, False)
+)
+def _O_jvp(forward_fn, params, model_state, samples, v, chunk_size):
+    @partial(scanmap, scan_fun=scan_append, argnums=3)
+    def __O_jvp(forward_fn, params, model_state, samples, v):
         # TODO apply the transpose of sum_inplace (allreduce) to the arg v here
         # in order to get correct transposition with JAX sharding
-        _, res = jax.jvp(lambda p: forward_fn(p, samples), (params,), (v,))
+        _, res = jax.jvp(lambda p: forward_fn(p, model_state, samples), (params,), (v,))
         return res
 
     samples, unchunk_fn = chunk(samples, chunk_size)
-    res = __O_jvp(forward_fn, params, samples, v)
+    res = __O_jvp(forward_fn, params, model_state, samples, v)
     return unchunk_fn(res)
 
 
 @partial(
     sharding_decorator,
-    sharded_args_tree=(False, False, True, True, False),
+    sharded_args_tree=(False, False, False, True, True, False),
     reduction_op_tree=jax.lax.psum,
 )
-def _O_vjp(forward_fn, params, samples, w, chunk_size):
-    @partial(scanmap, scan_fun=scan_reduce, argnums=(2, 3))
-    def __O_vjp(forward_fn, params, samples, w):
-        _, vjp_fun = jax.vjp(forward_fn, params, samples)
-        res, _ = vjp_fun(w)
-        return res
+def _O_vjp(forward_fn, params, model_state, samples, w, chunk_size):
+    @partial(scanmap, scan_fun=scan_reduce, argnums=(3, 4))
+    def __O_vjp(forward_fn, params, model_state, samples, w):
+        _, vjp_fun = jax.vjp(forward_fn, params, model_state, samples)
+        vjp_pars, _, _ = vjp_fun(w)
+        return vjp_pars
 
     samples, _ = chunk(samples, chunk_size)
     w, _ = chunk(w, chunk_size)
-    res = __O_vjp(forward_fn, params, samples, w)
+    res = __O_vjp(forward_fn, params, model_state, samples, w)
     return res
 
 
-def _OH_w(forward_fn, params, samples, w, chunk_size):
-    return tree_conj(_O_vjp(forward_fn, params, samples, w.conjugate(), chunk_size))
+def _OH_w(forward_fn, params, model_state, samples, w, chunk_size):
+    return tree_conj(
+        _O_vjp(forward_fn, params, model_state, samples, w.conjugate(), chunk_size)
+    )
 
 
-def _Odagger_DeltaO_v(forward_fn, params, samples, v, chunk_size, pdf=None):
-    w = _O_jvp(forward_fn, params, samples, v, chunk_size)
+def _Odagger_DeltaO_v(
+    forward_fn, params, model_state, samples, v, chunk_size, pdf=None
+):
+    w = _O_jvp(forward_fn, params, model_state, samples, v, chunk_size)
     if pdf is None:
         w = w * (1.0 / samples.shape[0])
         w = subtract_mean(w)  # w/ JAX sharding
     else:
         w = pdf * (w - pdf @ w)
-    res = _OH_w(forward_fn, params, samples, w, chunk_size)
+    res = _OH_w(forward_fn, params, model_state, samples, w, chunk_size)
     return res
 
 
-# @partial(jax.jit, static_argnums=1)
-def _mat_vec_chunked(forward_fn, params, samples, v, diag_shift, chunk_size, pdf=None):
+def _mat_vec_chunked(
+    forward_fn, params, model_state, samples, v, diag_shift, chunk_size, pdf=None
+):
     assert samples.ndim == 2  # require flat samples, axis 0 can be sharded
-    res = _Odagger_DeltaO_v(forward_fn, params, samples, v, chunk_size, pdf)
+    res = _Odagger_DeltaO_v(
+        forward_fn, params, model_state, samples, v, chunk_size, pdf
+    )
     return tree_axpy(diag_shift, v, res)
 
 
 def _mat_vec_chunked_transposable(
-    forward_fn, params, samples, v, diag_shift, chunk_size, pdf=None
+    forward_fn, params, model_state, samples, v, diag_shift, chunk_size, pdf=None
 ):
-    extra_args = (params, samples, diag_shift, pdf)
+    extra_args = (params, model_state, samples, diag_shift, pdf)
 
     def _mv(extra_args, x):
-        params, samples, diag_shift, pdf = extra_args
+        params, model_state, samples, diag_shift, pdf = extra_args
         return _mat_vec_chunked(
-            forward_fn, params, samples, x, diag_shift, chunk_size, pdf
+            forward_fn, params, model_state, samples, x, diag_shift, chunk_size, pdf
         )
 
     def _mv_trans(extra_args, y):
         # the linear operator is hermitian
-        params, samples, diag_shift, pdf = extra_args
+        params, model_state, samples, diag_shift, pdf = extra_args
         return tree_conj(
             _mat_vec_chunked(
-                forward_fn, params, samples, tree_conj(y), diag_shift, chunk_size, pdf
+                forward_fn,
+                params,
+                model_state,
+                samples,
+                tree_conj(y),
+                diag_shift,
+                chunk_size,
+                pdf,
             )
         )
 
@@ -198,13 +214,13 @@ def mat_vec_chunked_factory(
         lambda v,δ : (S + δ I) v
     """
 
-    def fun(W, samples):
-        return forward_fn({"params": W, **model_state}, samples)
+    def fun(W, ms, samples):
+        return forward_fn({"params": W, **ms}, samples)
 
     return Partial(
         partial(_mat_vec_chunked_transposable, fun, chunk_size=chunk_size),
         params,
+        model_state,
         samples,
         pdf=pdf,
     )
-    # return Partial(lambda f, *args: jax.jit(f)(*args), Partial(partial(_mat_vec_chunked, fun), params, samples))

{netket-3.19.0 → netket-3.19.2}/netket/sampler/parallel_tempering.py

@@ -42,18 +42,38 @@ class ParallelTemperingSamplerState(MetropolisSamplerState):
     Contains the usual quantities, as well as statistics about the paralel tempering.
     """
 
-    beta: jnp.ndarray = None
+    beta: jnp.ndarray = struct.field(
+        sharded=struct.ShardedFieldSpec(
+            sharded=True, deserialization_function="relaxed-ignore-errors"
+        )
+    )
     """The inverse temperatures of the different chains."""
 
-    n_accepted_per_beta: jnp.ndarray = None
+    n_accepted_per_beta: jnp.ndarray = struct.field(
+        sharded=struct.ShardedFieldSpec(
+            sharded=True, deserialization_function="relaxed-ignore-errors"
+        )
+    )
     """Total number of moves accepted per beta across all JAX processes since the last reset."""
-    beta_0_index: jnp.ndarray = None
+    beta_0_index: jnp.ndarray = struct.field(
+        sharded=struct.ShardedFieldSpec(
+            sharded=True, deserialization_function="relaxed-ignore-errors"
+        )
+    )
     r"""Index of the position of the chain with :math:`\\beta=1`."""
-    beta_position: jnp.ndarray = None
+    beta_position: jnp.ndarray = struct.field(
+        sharded=struct.ShardedFieldSpec(
+            sharded=True, deserialization_function="relaxed-ignore-errors"
+        )
+    )
     r"""Averaged position of :math:`\\beta=1`."""
-    beta_diffusion: jnp.ndarray = None
+    beta_diffusion: jnp.ndarray = struct.field(
+        sharded=struct.ShardedFieldSpec(
+            sharded=True, deserialization_function="relaxed-ignore-errors"
+        )
+    )
     """Average variance of the position of :math:`\\beta = 1`."""
-    exchange_steps: int = 0
+    exchange_steps: int = struct.field(default_factory=lambda: jnp.zeros((), dtype=int))
     """Number of exchanges between the different temperatures."""
 
     def __init__(

{netket-3.19.0 → netket-3.19.2}/netket/vqs/mc/mc_state/expect_forces_chunked.py

@@ -128,13 +128,14 @@ def forces_expect_hermitian_chunked(
     # mutable state (if it's there)
     if mutable is False:
         vjp_fun_chunked = nkjax.vjp_chunked(
-            lambda w, σ: model_apply_fun({"params": w, **model_state}, σ),
+            lambda w, ms, σ: model_apply_fun({"params": w, **ms}, σ),
             parameters,
+            model_state,
             σ,
             conjugate=True,
             chunk_size=chunk_size,
-            chunk_argnums=1,
-            nondiff_argnums=1,
+            chunk_argnums=2,
+            nondiff_argnums=(1, 2),
         )
         new_model_state = None
     else: