NetKet 3.19.1__tar.gz → 3.19.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.19.1
+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/#
@@ -22,7 +22,7 @@ Requires-Dist: einops>=0.8.1
 Requires-Dist: equinox>=0.11.3
 Requires-Dist: flax>=0.10.2
 Requires-Dist: igraph>=0.10.0
-Requires-Dist: jax>=0.7.0
+Requires-Dist: jax<0.6,>=0.4.35
 Requires-Dist: numba>=0.59
 Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.2.2

{netket-3.19.1 → netket-3.19.2}/netket/_src/ngd/sr.py

@@ -3,13 +3,11 @@ from functools import partial
 
 import jax
 import jax.numpy as jnp
-from jax.sharding import NamedSharding, PartitionSpec as P
+from jax.sharding import PositionalSharding
 
 from netket import jax as nkjax
 from netket.utils.types import Array
 
-from netket import config
-
 from netket._src.ngd.kwargs import ensure_accepts_kwargs
 
 
@@ -65,9 +63,8 @@ def _compute_sr_update(
         num_p = updates.shape[-1] // 2
         updates = updates[:num_p] + 1j * updates[num_p:]
 
-    if config.netket_experimental_sharding:
-        updates = jax.lax.with_sharding_constraint(
-            updates, NamedSharding(jax.sharding.get_abstract_mesh(), P())
-        )
+    updates = jax.lax.with_sharding_constraint(
+        updates, PositionalSharding(jax.devices()).replicate()
+    )
 
     return updates, old_updates, info

{netket-3.19.1 → netket-3.19.2}/netket/_src/ngd/srt.py

@@ -3,7 +3,7 @@ from functools import partial
 
 import jax
 import jax.numpy as jnp
-from jax.sharding import NamedSharding, PartitionSpec as P
+from jax.sharding import PositionalSharding
 
 from netket import jax as nkjax
 from netket import config
@@ -40,10 +40,10 @@ def _compute_srt_update(
         nkjax.sharding.pad_axis_for_sharding(O_LT, axis=1, padding_value=0.0)
         O_LT = jax.lax.with_sharding_constraint(
             O_LT,
-            NamedSharding(jax.sharding.get_abstract_mesh(), P("S", None)),
+            PositionalSharding(jax.devices()).reshape(1, -1),
         )
         dv = jax.lax.with_sharding_constraint(
-            dv, NamedSharding(jax.sharding.get_abstract_mesh(), P())
+            dv, PositionalSharding(jax.devices()).replicate()
         )
 
     # This does the contraction (ns, #np) x (#np, ns) -> (ns, ns).
@@ -85,9 +85,9 @@ def _compute_srt_update(
         updates = updates[:num_p] + 1j * updates[num_p:]
 
     if config.netket_experimental_sharding:
-        out_shardings = NamedSharding(
-            jax.sharding.get_abstract_mesh(), P(*(None,) * updates.ndim)
-        )
+        out_shardings = (
+            PositionalSharding(jax.devices()).replicate().reshape((1,) * updates.ndim)
+        ).replicate()
         updates = jax.lax.with_sharding_constraint(updates, out_shardings)
 
     return updates, old_updates, info

{netket-3.19.1 → netket-3.19.2}/netket/_src/ngd/srt_onthefly.py

@@ -7,13 +7,15 @@ import jax
 import jax.numpy as jnp
 from jax.tree_util import tree_map
 
-from jax.sharding import NamedSharding, PartitionSpec as P
+from jax.experimental.shard_map import shard_map
+from jax.sharding import PartitionSpec as P, PositionalSharding
 
 from netket import jax as nkjax
 from netket import config
 from netket.jax._jacobian.default_mode import JacobianMode
 from netket.utils import timing
 from netket.utils.types import Array
+from netket.utils.version_check import module_version
 
 from netket.jax import _ntk as nt
 
@@ -102,10 +104,10 @@ def srt_onthefly(
     all_samples = samples
     if config.netket_experimental_sharding:
         samples = jax.lax.with_sharding_constraint(
-            samples, NamedSharding(jax.sharding.get_abstract_mesh(), P("S", None))
+            samples, PositionalSharding(jax.devices()).reshape(-1, 1)
         )
         all_samples = jax.lax.with_sharding_constraint(
-            samples, NamedSharding(jax.sharding.get_abstract_mesh(), P())
+            samples, PositionalSharding(jax.devices()).replicate()
         )
 
     _jacobian_contraction = nt.empirical_ntk_by_jacobian(
@@ -115,8 +117,6 @@ def srt_onthefly(
     )
 
     def jacobian_contraction(samples, all_samples, parameters_real):
-        if config.netket_experimental_sharding:
-            parameters_real = jax.lax.pvary(parameters_real, "S")
         if chunk_size is None:
             # STRUCTURED_DERIVATIVES returns a complex array, but the imaginary part is zero
             # shape [N_mc/p.size, N_mc, 2, 2]
@@ -136,20 +136,27 @@ def srt_onthefly(
 
     # If we are sharding, use shard_map manually
     if config.netket_experimental_sharding:
-        mesh = jax.sharding.get_abstract_mesh()
+        mesh = jax.make_mesh((jax.device_count(),), ("i",), devices=jax.devices())
         # SAMPLES, ALL_SAMPLES PARAMETERS_REAL
-        in_specs = (P("S", None), P(), P())
-        out_specs = P("S", None)
+        in_specs = (P("i", None), P(), P())
+        out_specs = P("i", None, None, None)
 
         # By default, I'm not sure whether the jacobian_contraction of NeuralTangents
         # Is correctly automatically sharded across devices. So we force it to be
         # sharded with shard map to be sure
 
-        jacobian_contraction = jax.shard_map(
+        # check rep:
+        check_rep = module_version("jax") < (0, 4, 38)
+        # shard_map is broken between 0.4.38 and (as of 25 march 2025) 0.5.3.
+        # We assume any version after 0.4.38 'has a bug' that shows up as
+        # None is not Iterable
+        # it's a bug in check_rep, so we disable it in this case
+        jacobian_contraction = shard_map(
             jacobian_contraction,
             mesh=mesh,
             in_specs=in_specs,
             out_specs=out_specs,
+            check_rep=check_rep,
         )
 
     # This disables the nkjax.sharding_decorator in here, which might appear
@@ -158,12 +165,10 @@ def srt_onthefly(
         ntk_local = jacobian_contraction(samples, all_samples, parameters_real).real
 
     # shape [N_mc, N_mc, 2, 2] or [N_mc, N_mc]
-    if config.netket_experimental_sharding:
-        ntk = jax.lax.with_sharding_constraint(
-            ntk_local, NamedSharding(jax.sharding.get_abstract_mesh(), P())
-        )
-    else:
-        ntk = ntk_local
+    ntk = jax.lax.with_sharding_constraint(
+        ntk_local,
+        PositionalSharding(jax.devices()).replicate(),
+    )
     if mode == "complex":
         # shape [2*N_mc, 2*N_mc] checked with direct calculation of J^T J
         ntk = rearrange(ntk, "i j z w -> (i z) (j w)")
@@ -210,14 +215,12 @@ def srt_onthefly(
     if mode == "complex":
         aus_vector = aus_vector.reshape(-1, 2)
     # shape [N_mc // p.size,2]
-    if config.netket_experimental_sharding:
-        aus_vector = jax.lax.with_sharding_constraint(
-            aus_vector,
-            NamedSharding(
-                jax.sharding.get_abstract_mesh(),
-                P("S", *(None,) * (aus_vector.ndim - 1)),
-            ),
-        )
+    aus_vector = jax.lax.with_sharding_constraint(
+        aus_vector,
+        PositionalSharding(jax.devices()).reshape(
+            -1, *(1 for _ in range(aus_vector.ndim - 1))
+        ),
+    )
 
     # _, vjp_fun = jax.vjp(f, parameters_real)
     vjp_fun = nkjax.vjp_chunked(

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

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.19.1'
-__version_tuple__ = version_tuple = (3, 19, 1)
+__version__ = version = '3.19.2'
+__version_tuple__ = version_tuple = (3, 19, 2)
 
 __commit_id__ = commit_id = None

{netket-3.19.1 → netket-3.19.2}/netket/driver/abstract_variational_driver.py

@@ -22,7 +22,6 @@ from functools import partial
 from tqdm.auto import tqdm
 
 import jax
-from jax.sharding import NamedSharding, PartitionSpec as P
 
 from netket import config
 from netket.logging import AbstractLog, JsonLog
@@ -201,7 +200,7 @@ class AbstractVariationalDriver(abc.ABC):
             if config.netket_experimental_sharding:
                 self._optimizer_state = jax.lax.with_sharding_constraint(
                     self._optimizer_state,
-                    NamedSharding(jax.sharding.get_abstract_mesh(), P()),
+                    jax.sharding.PositionalSharding(jax.devices()).replicate(),
                 )
 
     @property
@@ -435,7 +434,7 @@ def apply_gradient(optimizer_fun, optimizer_state, dp, params):
     new_params = optax.apply_updates(params, updates)
 
     if config.netket_experimental_sharding:
-        sharding = NamedSharding(jax.sharding.get_abstract_mesh(), P())
+        sharding = jax.sharding.PositionalSharding(jax.devices()).replicate()
         new_optimizer_state = jax.lax.with_sharding_constraint(
             new_optimizer_state, sharding
         )

{netket-3.19.1 → netket-3.19.2}/netket/errors.py

@@ -1116,9 +1116,14 @@ def concrete_or_error(force, value, error_class, *args, **kwargs):
       *args: any additional argument and keyword argument to pass to the custom
         error type constructor.
     """
+    import jax
     from jax.core import concrete_or_error
 
-    from jax.errors import ConcretizationTypeError
+    # TODO: remove once we require jax >= 0.5.0
+    if hasattr(jax.errors, "ConcretizationTypeError"):
+        from jax.errors import ConcretizationTypeError
+    else:
+        from jax.core import ConcretizationTypeError
 
     try:
         return concrete_or_error(

{netket-3.19.1 → netket-3.19.2}/netket/experimental/qsr/driver.py

@@ -134,8 +134,7 @@ class QSR(AbstractVariationalDriver):
         if not isinstance(training_data, RawQuantumDataset):
             training_data = RawQuantumDataset(training_data)
 
-        seed = nkjax.PRNGKey(seed)
-        self._rng = np.random.default_rng(np.asarray(jax.random.key_data(seed)))
+        self._rng = np.random.default_rng(np.asarray(nkjax.PRNGKey(seed)))
 
         # mixed states
         self.mixed_states = variational_state.__class__.__name__ in ["MCMixedState"]

{netket-3.19.1 → netket-3.19.2}/netket/hilbert/discrete_hilbert.py

@@ -161,7 +161,7 @@ class DiscreteHilbert(AbstractHilbert):
 
         # equinox.error_if is broken under shard_map.
         # If we are using shard map, we skip this check
-        if sharding._get_SHARD_MAP_STACK_LEVEL() == 0 and jax.device_count() == 1:
+        if sharding.SHARD_MAP_STACK_LEVEL == 0 and jax.device_count() == 1:
             numbers = error_if(
                 numbers,
                 (numbers >= self.n_states).any() | (numbers < 0).any(),

{netket-3.19.1 → netket-3.19.2}/netket/hilbert/homogeneous.py

@@ -264,7 +264,7 @@ class HomogeneousHilbert(DiscreteHilbert):
 
             # equinox.error_if is broken under shard_map.
             # If we are using shard map, we skip this check
-            if sharding._get_SHARD_MAP_STACK_LEVEL() == 0 and jax.device_count() == 1:
+            if sharding.SHARD_MAP_STACK_LEVEL == 0 and jax.device_count() == 1:
                 states = error_if(
                     states,
                     (states < start).any() | (states >= end).any(),
@@ -272,7 +272,7 @@ class HomogeneousHilbert(DiscreteHilbert):
                 )
 
         if self.constrained:
-            if sharding._get_SHARD_MAP_STACK_LEVEL() == 0 and jax.device_count() == 1:
+            if sharding.SHARD_MAP_STACK_LEVEL == 0 and jax.device_count() == 1:
                 states = error_if(
                     states,
                     ~self.constraint(states).all(),

{netket-3.19.1 → netket-3.19.2}/netket/hilbert/random/base.py

@@ -14,7 +14,7 @@
 
 from textwrap import dedent
 from typing import Union
-from functools import partial
+
 
 import jax
 import numpy as np
@@ -32,7 +32,7 @@ def random_state(hilb, key, *, size=None, dtype=np.float32):
 
     Args:
         hilb: The hilbert space
-        key: The rng key generated by :func:`jax.random.key`.
+        key: The Jax :class:`jax.random.KeyArray` generated by :func:`jax.random.PRNGKey`.
         size: If provided, returns a batch of configurations of the form `(size, #)` if
             size is an integer or `(*size, #)` if it is a tuple and where `#` is the Hilbert
             space size. By default, a single random configuration with shape `(#,)` is
@@ -115,7 +115,6 @@ def flip_state_scalar(hilb, key, state, indx):
 
 
 @dispatch
-@partial(jax.jit, static_argnames=("hilb",))
 def flip_state_batch(hilb, key, states, indxs):
     keys = jax.random.split(key, states.shape[0])
     res = jax.vmap(flip_state_scalar, in_axes=(None, 0, 0, 0), out_axes=0)(

{netket-3.19.1 → netket-3.19.2}/netket/jax/_chunk_utils.py

@@ -2,9 +2,7 @@ from functools import partial
 
 import jax
 import jax.numpy as jnp
-from netket.utils.iterators import safe_map
-
-safe_zip = partial(zip, strict=True)
+from jax.util import safe_zip, safe_map
 
 
 def _tree_map_multi_out(f, tree, *rest, is_leaf=None):

{netket-3.19.1 → netket-3.19.2}/netket/jax/_grad.py

@@ -19,8 +19,12 @@ import operator
 
 import jax
 import jax.numpy as jnp
+from jax.util import safe_map
+
 from ._utils_tree import tree_leaf_iscomplex, eval_shape
 
+map = safe_map
+
 
 def _ensure_index(x: Any) -> int | tuple[int, ...]:
     """Ensure x is either an index or a tuple of indices."""

{netket-3.19.1 → netket-3.19.2}/netket/jax/_jacobian/jacobian_pytree.py

@@ -19,6 +19,7 @@ import jax
 import jax.flatten_util
 import jax.numpy as jnp
 
+import numpy as np
 
 from netket import jax as nkjax
 from netket.utils import wrap_to_support_scalar
@@ -41,7 +42,7 @@ def jacobian_real_holo(forward_fn: Callable, params: PyTree, samples: Array) ->
     y, vjp_fun = jax.vjp(
         lambda pars: wrap_to_support_scalar(forward_fn)(pars, samples), params
     )
-    (res,) = vjp_fun(jax.lax.full_like(y, 1.0))
+    (res,) = vjp_fun(np.array(1.0, dtype=jnp.result_type(y)))
     return res
 
 
@@ -70,8 +71,8 @@ def _jacobian_cplx(
     if not jnp.issubdtype(jnp.result_type(y), jnp.complexfloating):
         raise TypeError("Cannot build the complex jacobian for a real-valued function.")
 
-    (gr,) = vjp_fun(jax.lax.full_like(y, 1))
-    (gi,) = vjp_fun(jax.lax.full_like(y, -1j))
+    (gr,) = vjp_fun(np.array(1.0, dtype=jnp.result_type(y)))
+    (gi,) = vjp_fun(np.array(-1.0j, dtype=jnp.result_type(y)))
     return _build_fn(gr, gi)
 
 

{netket-3.19.1 → netket-3.19.2}/netket/jax/_jacobian/logic.py

@@ -27,6 +27,7 @@ from netket.jax import (
     tree_to_real,
     vmap_chunked,
 )
+from netket.jax.sharding import sharding_decorator
 
 from . import jacobian_dense
 from . import jacobian_pytree
@@ -42,7 +43,6 @@ from . import jacobian_pytree
         "center",
         "dense",
         "_sqrt_rescale",
-        "_axis_0_is_sharded",
     ),
 )
 def jacobian(
@@ -354,6 +354,21 @@ def jacobian(
         chunk_size=chunk_size,
         axis_0_is_sharded=_axis_0_is_sharded,
     )  # see below
+    # vmap_chunked, as of 5/11/2024 (3.14.3) only passses the function through
+    # a sharding_decorator if it is chunked, and not if it's not. This creates
+    # a different behaviour in the two cases. To homogenize it, we do it at the
+    # level of nkjax.jacobian.
+    # This fixes computing the jacobian of logpsi with sharded functions like
+    # for example from get_conn_padded.
+
+    # TODO: Maybe remove this and simply always sharding_decorator inside of vmap_chunked.
+    # We do sharding decorator here to account for cases where the wavefunction contains
+    if _axis_0_is_sharded is True:
+        jacobian_fun = sharding_decorator(
+            jacobian_fun,
+            (False, False, True),
+            reduction_op_tree=False,
+        )
 
     jacobians = jacobian_fun(Partial(f), params, samples)
 

{netket-3.19.1 → netket-3.19.2}/netket/jax/_ntk/logic.py

@@ -28,20 +28,28 @@ For details, please see "`Fast Finite Width Neural Tangent Kernel
 <https://arxiv.org/abs/2206.08720>`_".
 """
 
+import operator
 from typing import Any, Protocol, TypeVar, Union
 from collections.abc import Callable, Iterable
 
-from functools import partial
-import operator
 
-import jax
-import jax.numpy as jnp
+from jax import eval_shape
+from jax import jacobian
+from jax import lax
+from jax import vmap
+
 from jax.core import ShapedArray
+
+
 from jax.interpreters.ad import UndefinedPrimal
 
-from . import utils
+import jax.numpy as jnp
 
-safe_zip = partial(zip, strict=True)
+from jax.tree_util import tree_map
+from jax.tree_util import tree_reduce
+from jax.util import safe_zip as zip
+
+from . import utils
 
 PyTree = Any
 _VMapAxis = PyTree | None
@@ -223,7 +231,7 @@ def empirical_ntk_by_jacobian(
             contract_axes = _trace_axes + param_axes
             return _dot_general(x, y, contract_axes, _diagonal_axes) / size
 
-        return jax.tree.reduce(operator.add, jax.tree.map(contract, j1, j2))
+        return tree_reduce(operator.add, tree_map(contract, j1, j2))
 
     def ntk_fn(
         x1: PyTree, x2: PyTree | None, params: PyTree, **apply_fn_kwargs
@@ -264,16 +272,16 @@ def empirical_ntk_by_jacobian(
         def j_fn(x, *args):
             _kwargs = {k: v for k, v in zip(keys, args)}
             fx = _get_f_params(f, x, x_axis, fx_axis, kw_axes, **_kwargs)
-            jx = jax.jacobian(fx)(params)
+            jx = jacobian(fx)(params)
             return jx
 
         if not utils.all_none(x_axis) or not utils.all_none(kw_axes):
             in_axes = [x_axis] + [kw_axes[k] if k in kw_axes else None for k in keys]
-            j_fn = jax.vmap(j_fn, in_axes=in_axes, out_axes=fx_axis)
+            j_fn = vmap(j_fn, in_axes=in_axes, out_axes=fx_axis)
 
         j1 = j_fn(x1, *args1)
         j2 = j_fn(x2, *args2) if not utils.all_none(x2) else j1
-        ntk = jax.tree.map(sum_and_contract, fx1, j1, j2)
+        ntk = tree_map(sum_and_contract, fx1, j1, j2)
         return ntk
 
     return ntk_fn
@@ -305,13 +313,13 @@ def _get_f_params(
 
 
 def _ndim(x: PyTree) -> PyTree:
-    return jax.tree.map(lambda x: x.ndim, x)
+    return tree_map(lambda x: x.ndim, x)
 
 
 def _mod(x: PyTree | None, y: PyTree) -> PyTree:
     if x is None:
         return None
-    return jax.tree.map(operator.mod, x, y)
+    return tree_map(operator.mod, x, y)
 
 
 def _squeeze(x: PyTree, axis: PyTree | None) -> PyTree:
@@ -340,7 +348,7 @@ def _squeeze(x: PyTree, axis: PyTree | None) -> PyTree:
 
         return jnp.squeeze(x, non_zero_axes)
 
-    return jax.tree.map(squeeze, x, axis)
+    return tree_map(squeeze, x, axis)
 
 
 def _expand_dims_array(x: _ArrayOrShape, axis: int) -> _ArrayOrShape:
@@ -348,7 +356,7 @@ def _expand_dims_array(x: _ArrayOrShape, axis: int) -> _ArrayOrShape:
         return jnp.expand_dims(x, axis)
 
     if isinstance(x, ShapedArray):
-        return jax.eval_shape(expand, x)
+        return eval_shape(expand, x)
 
     if isinstance(x, jnp.ndarray):
         return expand(x)
@@ -361,7 +369,7 @@ def _expand_dims(
 ) -> PyTree | None:
     if axis is None or x is None or isinstance(x, UndefinedPrimal):
         return x
-    return jax.tree.map(_expand_dims_array, x, axis)
+    return tree_map(_expand_dims_array, x, axis)
 
 
 def _get_args(
@@ -374,8 +382,8 @@ def _get_args(
 ):
     kwargs1, kwargs2 = utils.split_kwargs(apply_fn_kwargs, x1, x2)
 
-    fx1 = jax.eval_shape(f, params, x1, **kwargs1)
-    fx2 = fx1 if utils.all_none(x2) else jax.eval_shape(f, params, x2, **kwargs2)
+    fx1 = eval_shape(f, params, x1, **kwargs1)
+    fx2 = fx1 if utils.all_none(x2) else eval_shape(f, params, x2, **kwargs2)
 
     x_axis, fx_axis, kw_axes = _canonicalize_axes(vmap_axes, x1, fx1, **kwargs1)
 
@@ -394,13 +402,13 @@ def _canonicalize_axes(
         x_axis, fx_axis, kw_axes = vmap_axes, vmap_axes, {}
 
     if isinstance(x_axis, int):
-        x_axis = jax.tree.map(lambda _: x_axis, x)
+        x_axis = tree_map(lambda _: x_axis, x)
 
     if isinstance(fx_axis, int):
-        fx_axis = jax.tree.map(lambda _: fx_axis, fx)
+        fx_axis = tree_map(lambda _: fx_axis, fx)
 
     if isinstance(kw_axes, int):
-        kw_axes = jax.tree.map(lambda _: kw_axes, kwargs)
+        kw_axes = tree_map(lambda _: kw_axes, kwargs)
 
     x_axis = _mod(x_axis, _ndim(x))
     fx_axis = _mod(fx_axis, _ndim(fx))
@@ -459,7 +467,7 @@ def _dot_general(
 
     dimension_numbers = ((contracting_dims, contracting_dims), (batch_dims, batch_dims))
 
-    prod = jax.lax.dot_general(lhs, rhs, dimension_numbers, precision)
+    prod = lax.dot_general(lhs, rhs, dimension_numbers, precision)
     prod = utils.zip_axes(prod, n_batch_dims)
 
     res_batch_dims = _get_res_batch_dims(contracting_dims, batch_dims)

{netket-3.19.1 → netket-3.19.2}/netket/jax/_scanmap.py

@@ -10,7 +10,9 @@ from jax.extend import linear_util as lu
 from ._utils_tree import compose
 
 _tree_add = partial(jax.tree_util.tree_map, jax.lax.add)
-_tree_zeros_like = partial(jax.tree_util.tree_map, jnp.zeros_like)
+_tree_zeros_like = partial(
+    jax.tree_util.tree_map, lambda x: jnp.zeros(x.shape, dtype=x.dtype)
+)
 
 
 # TODO put it somewhere
@@ -88,8 +90,7 @@ def scan_append_reduce(f, x, append_cond, op=_tree_add, zero_fun=_tree_zeros_lik
     _get_op_part = partial(_multimap, lambda c, x: x if not c else None, append_cond)
     _tree_select = partial(_multimap, lambda c, t1, t2: t1 if c else t2, append_cond)
 
-    dummy_y = jax.eval_shape(f_flat, x0)
-    carry_init = True, _get_op_part(zero_fun(dummy_y))
+    carry_init = True, _get_op_part(zero_fun(jax.eval_shape(f_flat, x0)))
 
     def f_(carry, x):
         is_first, y_carry = carry

{netket-3.19.1 → netket-3.19.2}/netket/jax/_utils_random.py

@@ -15,7 +15,6 @@
 
 import jax
 import jax.numpy as jnp
-from jax.sharding import NamedSharding, PartitionSpec as P
 
 from netket.utils import random_seed, config
 from netket.utils.types import PRNGKeyT, SeedT
@@ -51,28 +50,16 @@ def PRNGKey(seed: SeedT | None = None, *, root: int = 0) -> PRNGKeyT:
                 ).item()
             )
 
-        key = jax.random.key(seed)
-    elif isinstance(seed, jax.Array):
-        if jnp.issubdtype(seed.dtype, jax.dtypes.prng_key):
-            # new-style keys
-            key = seed
-        elif jnp.issubdtype(seed.dtype, jax.numpy.unsignedinteger):
-            # old-style PRNGKey
-            key = seed
-        else:
-            raise TypeError(
-                f"Expected seed to be an integer or a PRNGKey, got {type(seed)}, {seed.dtype} : {seed}"
-            )
+        key = jax.random.PRNGKey(seed)
     else:
-        raise TypeError(f"unsupported type {type(seed)}")
-
-    mesh = jax.sharding.get_abstract_mesh()
-    if not mesh.empty:
-        if len(mesh.explicit_axes) >= 0:
-            key = jax.sharding.reshard(key, P())
-        else:
-            sharding = NamedSharding(mesh, P())
-            key = jax.lax.with_sharding_constraint(key, sharding)
+        key = seed
+
+    if config.netket_experimental_sharding:
+        key = jax.lax.with_sharding_constraint(
+            key, jax.sharding.PositionalSharding(jax.devices()).replicate()
+        )
+    else:  # type: ignore[attr-defined]
+        key = _bcast_key(key, root=root)
     return key
 
 

{netket-3.19.1 → netket-3.19.2}/netket/jax/_vjp_chunked.py

@@ -100,7 +100,7 @@ def __vjp_fun_chunked(
             if c:
                 return _tree_unchunk(l, r)
             else:
-                return jax.tree_util.tree_map(jax.lax.add, l, r)
+                return jax.tree_map(jax.lax.add, l, r)
 
         return jax.tree_util.tree_map(_f, append_cond, res_chunk, res_rest)
     elif n_chunks > 0:
@@ -316,12 +316,6 @@ def vjp_chunked(
             lambda c: jax.lax.psum if c else False, out_args
         )
 
-        # jax will psum in every vjp, but we want to do one by hand at the end
-        # so we need to set the sharded args we differentiate as varying
-        # see https://github.com/jax-ml/jax/discussions/29608
-        pvary_argnums = tuple(set(range(len(primals))).difference(nondiff_argnums))
-        pvary_args = tuple(i in pvary_argnums for i in range(len(primals)))
-
         # check the chunk_size is not larger than the shard per device
         chunk_size = sharding_decorator(
             partial(check_chunk_size, chunk_argnums, chunk_size),
@@ -353,7 +347,6 @@ def vjp_chunked(
                     _vjpc,
                     sharded_args_tree=(sharded_args, True),
                     reduction_op_tree=reduction_op_tree,
-                    pvary_args_tree=(pvary_args, False),
                 ),
                 primals,
             )

{netket-3.19.1 → netket-3.19.2}/netket/jax/_vmap_chunked.py

@@ -42,15 +42,7 @@ def _eval_fun_in_chunks_sharding(vmapped_fun, chunk_size, argnums, *args, **kwar
     # by computing the vmapped_fun in chunks on every shard (which sits on a separate device)
     sharded_args_tree = tuple(i in argnums for i, a in enumerate(args))
     f = HashablePartial(_eval_fun_in_chunks, vmapped_fun, chunk_size, argnums, **kwargs)
-
-    # if the vmapped_fun e.g. does a vjp we need to make the params pvary here
-    # to avoid it emitting an unwanted psum
-    pvary_argnums = tuple(set(range(len(args))).difference(argnums))
-    pvary_args_tree = tuple(i in pvary_argnums for i in range(len(args)))
-
-    return sharding_decorator(f, sharded_args_tree, pvary_args_tree=pvary_args_tree)(
-        *args
-    )
+    return sharding_decorator(f, sharded_args_tree)(*args)
 
 
 def _chunk_vmapped_function(