NetKet 3.11.3__tar.gz → 3.11.4__tar.gz

{netket-3.11.3 → netket-3.11.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.11.3
+Version: 3.11.4
 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.11.3 → netket-3.11.4}/netket/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.11.3'
-__version_tuple__ = version_tuple = (3, 11, 3)
+__version__ = version = '3.11.4'
+__version_tuple__ = version_tuple = (3, 11, 4)

{netket-3.11.3 → netket-3.11.4}/netket/driver/abstract_variational_driver.py

@@ -376,7 +376,7 @@ class AbstractVariationalDriver(abc.ABC):
 
         # Do not unpack operators, even if they are pytrees!
         # this is necessary to support jax operators.
-        return jax.tree_map(
+        return jax.tree_util.tree_map(
             self._estimate_stats,
             observables,
             is_leaf=lambda x: isinstance(x, AbstractObservable),

{netket-3.11.3 → netket-3.11.4}/netket/errors.py

@@ -587,7 +587,7 @@ class RealQGTComplexDomainError(NetketError):
        >>> _, vec = vstate.expect_and_grad(nk.operator.spin.sigmax(vstate.hilbert, 1))
        >>> G = nk.optimizer.qgt.QGTOnTheFly(vstate, holomorphic=False)
        >>>
-       >>> vec_real = jax.tree_map(lambda x: x.real, vec)
+       >>> vec_real = jax.tree.map(lambda x: x.real, vec)
        >>> sol = G@vec_real
 
    Or, if you used the QGT in a linear solver, try using:
@@ -601,7 +601,7 @@ class RealQGTComplexDomainError(NetketError):
        >>> _, vec = vstate.expect_and_grad(nk.operator.spin.sigmax(vstate.hilbert, 1))
        >>>
        >>> G = nk.optimizer.qgt.QGTOnTheFly(vstate, holomorphic=False)
-       >>> vec_real = jax.tree_map(lambda x: x.real, vec)
+       >>> vec_real = jax.tree.map(lambda x: x.real, vec)
        >>>
        >>> linear_solver = jax.scipy.sparse.linalg.cg
        >>> solution, info = G.solve(linear_solver, vec_real)
@@ -624,14 +624,14 @@ class RealQGTComplexDomainError(NetketError):
 
             .. code:: python
 
-               >>> vec_real = jax.tree_map(lambda x: x.real, vec)
+               >>> vec_real = jax.tree_util.tree_map(lambda x: x.real, vec)
                >>> G@vec_real
 
             If you used the QGT in a linear solver, try using:
 
             .. code:: python
 
-               >>> vec_real = jax.tree_map(lambda x: x.real, vec)
+               >>> vec_real = jax.tree_util.tree_map(lambda x: x.real, vec)
                >>> G.solve(linear_solver, vec_real)
 
             to fix this error.

{netket-3.11.3 → netket-3.11.4}/netket/experimental/driver/tdvp.py

@@ -162,7 +162,7 @@ def odefun_tdvp(  # noqa: F811
 
 @partial(jax.jit, static_argnums=(3, 4))
 def _map_parameters(forces, parameters, loss_grad_factor, propagation_type, state_T):
-    forces = jax.tree_map(
+    forces = jax.tree_util.tree_map(
         lambda x, target: loss_grad_factor * x,
         forces,
         parameters,

{netket-3.11.3 → netket-3.11.4}/netket/experimental/driver/tdvp_common.py

@@ -478,7 +478,7 @@ def odefun_host_callback(state, driver, *args, **kwargs):
     if config.netket_experimental_disable_ode_jit:
         return odefun(state, driver, *args, **kwargs)
 
-    result_shape = jax.tree_map(
+    result_shape = jax.tree_util.tree_map(
         lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
         state.parameters,
     )

{netket-3.11.3 → netket-3.11.4}/netket/experimental/driver/tdvp_schmitt.py

@@ -280,7 +280,7 @@ def odefun_schmitt(state: MCState, self: TDVPSchmitt, t, w, *, stage=0):  # noqa
 
 @partial(jax.jit, static_argnums=(3, 4))
 def _map_parameters(forces, parameters, loss_grad_factor, propagation_type, state_T):
-    forces = jax.tree_map(
+    forces = jax.tree_util.tree_map(
         lambda x, target: loss_grad_factor * x,
         forces,
         parameters,

{netket-3.11.3 → netket-3.11.4}/netket/experimental/driver/vmc_srt.py

@@ -210,7 +210,7 @@ class VMC_SRt(AbstractVariationalDriver):
         self.jacobian_mode = jacobian_mode
         self._linear_solver_fn = linear_solver_fn
 
-        self._params_structure = jax.tree_map(
+        self._params_structure = jax.tree_util.tree_map(
             lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), self.state.parameters
         )
         if not nkjax.tree_ishomogeneous(self._params_structure):

{netket-3.11.3 → netket-3.11.4}/netket/experimental/dynamics/_rk_solver_structures.py

@@ -104,7 +104,7 @@ def euclidean_norm(x: Union[PyTree, Array]):
         return jnp.sqrt(
             jax.tree_util.tree_reduce(
                 lambda x, y: x + y,
-                jax.tree_map(lambda x: jnp.sum(jnp.abs(x) ** 2), x),
+                jax.tree_util.tree_map(lambda x: jnp.sum(jnp.abs(x) ** 2), x),
             )
         )
 
@@ -119,7 +119,7 @@ def maximum_norm(x: Union[PyTree, Array]):
         return jnp.sqrt(
             jax.tree_util.tree_reduce(
                 jnp.maximum,
-                jax.tree_map(lambda x: jnp.max(jnp.abs(x)), x),
+                jax.tree_util.tree_map(lambda x: jnp.max(jnp.abs(x)), x),
             )
         )
 

{netket-3.11.3 → netket-3.11.4}/netket/experimental/dynamics/_rk_tableau.py

@@ -32,7 +32,7 @@ def expand_dim(tree: PyTree, sz: int):
     def _expand(x):
         return jnp.zeros((sz, *x.shape), dtype=x.dtype)
 
-    return jax.tree_map(_expand, tree)
+    return jax.tree_util.tree_map(_expand, tree)
 
 
 @dataclass
@@ -114,19 +114,19 @@ class TableauRKExplicit:
 
         k = expand_dim(y_t, self.stages)
         for l in range(self.stages):
-            dy_l = jax.tree_map(
+            dy_l = jax.tree_util.tree_map(
                 lambda k: jnp.tensordot(
                     jnp.asarray(self.a[l], dtype=k.dtype), k, axes=1
                 ),
                 k,
             )
-            y_l = jax.tree_map(
+            y_l = jax.tree_util.tree_map(
                 lambda y_t, dy_l: jnp.asarray(y_t + dt * dy_l, dtype=dy_l.dtype),
                 y_t,
                 dy_l,
             )
             k_l = f(times[l], y_l, stage=l)
-            k = jax.tree_map(lambda k, k_l: k.at[l].set(k_l), k, k_l)
+            k = jax.tree_util.tree_map(lambda k, k_l: k.at[l].set(k_l), k, k_l)
 
         return k
 
@@ -141,7 +141,7 @@ class TableauRKExplicit:
         k = self._compute_slopes(f, t, dt, y_t)
 
         b = self.b[0] if self.b.ndim == 2 else self.b
-        y_tp1 = jax.tree_map(
+        y_tp1 = jax.tree_util.tree_map(
             lambda y_t, k: y_t
             + jnp.asarray(dt, dtype=y_t.dtype)
             * jnp.tensordot(jnp.asarray(b, dtype=k.dtype), k, axes=1),
@@ -167,7 +167,7 @@ class TableauRKExplicit:
 
         k = self._compute_slopes(f, t, dt, y_t)
 
-        y_tp1 = jax.tree_map(
+        y_tp1 = jax.tree_util.tree_map(
             lambda y_t, k: y_t
             + jnp.asarray(dt, dtype=y_t.dtype)
             * jnp.tensordot(jnp.asarray(self.b[0], dtype=k.dtype), k, axes=1),
@@ -175,7 +175,7 @@ class TableauRKExplicit:
             k,
         )
         db = self.b[0] - self.b[1]
-        y_err = jax.tree_map(
+        y_err = jax.tree_util.tree_map(
             lambda k: jnp.asarray(dt, dtype=k.dtype)
             * jnp.tensordot(jnp.asarray(db, dtype=k.dtype), k, axes=1),
             k,

{netket-3.11.3 → netket-3.11.4}/netket/experimental/qsr/driver.py

@@ -268,7 +268,7 @@ class QSR(AbstractVariationalDriver):
 
             # gather gradient
             # grad <- grad - grad_cv + E[grad_cv]
-            self._grad_pos = jax.tree_map(
+            self._grad_pos = jax.tree_util.tree_map(
                 lambda x, y, Ey: x - y + Ey,
                 self._grad_pos,
                 self._grad_pos_cv,
@@ -281,7 +281,7 @@ class QSR(AbstractVariationalDriver):
 
         # restore the square in prob = |psi|^2
         if not self.mixed_states:
-            self._loss_grad = jax.tree_map(lambda x: x * 2.0, self._loss_grad)
+            self._loss_grad = jax.tree_util.tree_map(lambda x: x * 2.0, self._loss_grad)
 
         # If parameters are real, then take only real part of the gradient (if it's complex)
         self._loss_grad = tree_cast(self._loss_grad, self.state.parameters)

{netket-3.11.3 → netket-3.11.4}/netket/experimental/qsr/logic_helpers.py

@@ -49,7 +49,7 @@ def _avg_O(afun, pars, model_state, sigma):
     sigma = sigma.reshape((-1, sigma.shape[-1]))
     _, vjp = nkjax.vjp(lambda W: afun({"params": W, **model_state}, sigma), pars)
     (O_avg,) = vjp(jnp.ones(sigma.shape[0]) / sigma.shape[0])
-    return jax.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], O_avg)
+    return jax.tree_util.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], O_avg)
 
 
 @dispatch
@@ -84,7 +84,7 @@ def _avg_O_exact(hilbert: AbstractHilbert, afun, pars, model_state):
     psi_2 = jnp.abs(jnp.exp(afun({"params": pars, **model_state}, sigma))) ** 2
     psi_2 /= jnp.sum(psi_2)
     (O_avg,) = vjp(psi_2)
-    return jax.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], O_avg)
+    return jax.tree_util.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], O_avg)
 
 
 @dispatch
@@ -181,7 +181,7 @@ def _grad_local_value_rotated(log_psi, pars, model_state, sigma_p, mel, secs):
 
     (O_avg,) = vjp(jnp.ones_like(log_val_rotated) / log_val_rotated.size)
 
-    O_avg = jax.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], O_avg)
+    O_avg = jax.tree_util.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], O_avg)
 
     return log_val_rotated, O_avg
 

{netket-3.11.3 → netket-3.11.4}/netket/jax/_chunk_utils.py

@@ -4,7 +4,7 @@ from functools import partial
 
 def _treeify(f):
     def _f(x, *args, **kwargs):
-        return jax.tree_map(lambda y: f(y, *args, **kwargs), x)
+        return jax.tree_util.tree_map(lambda y: f(y, *args, **kwargs), x)
 
     return _f
 

{netket-3.11.3 → netket-3.11.4}/netket/jax/_grad.py

@@ -188,7 +188,9 @@ def value_and_grad(
                         )(*args, **kwargs)
 
                     out = out_r + 1j * out_j
-                    grad = jax.tree_map(lambda re, im: re + 1j * im, grad_r, grad_j)
+                    grad = jax.tree_util.tree_map(
+                        lambda re, im: re + 1j * im, grad_r, grad_j
+                    )
 
                     if has_aux:
                         return out, grad, aux

{netket-3.11.3 → netket-3.11.4}/netket/jax/_jacobian/jacobian_pytree.py

@@ -50,7 +50,7 @@ def _jacobian_cplx(
     params: PyTree,
     samples: Array,
     *,
-    _build_fn: Callable = partial(jax.tree_map, jax.lax.complex),
+    _build_fn: Callable = partial(jax.tree_util.tree_map, jax.lax.complex),
 ) -> PyTree:
     """Calculates Jacobian entries by vmapping grad.
     Assumes the function is R→C, backpropagates 1 and -1j
@@ -85,7 +85,7 @@ def stack_jacobian_tuple(ok_re_im):
         centered_oks_re_im : a tuple (ΔOᵣ, ΔOᵢ) of two PyTrees representing the real and imag part of ΔOⱼₖ
     """
     re, im = ok_re_im
-    return jax.tree_map(lambda re, im: jnp.stack([re, im], axis=0), re, im)
+    return jax.tree_util.tree_map(lambda re, im: jnp.stack([re, im], axis=0), re, im)
 
 
 jacobian_real_holo_fun = jacobian_real_holo

{netket-3.11.3 → netket-3.11.4}/netket/jax/_jacobian/logic.py

@@ -158,7 +158,7 @@ def jacobian(
     .. code:: python
 
       samples = samples.reshape(-1, samples.shape[-1])
-      parameters = jax.tree_map(lambda x: x.real, parameters)
+      parameters = jax.tree_util.tree_map(lambda x: x.real, parameters)
       O_k = jax.jacrev(lambda pars: logpsi(pars, samples).real, parameters)
 
     The jacobian that is returned is a PyTree with the same shape
@@ -201,7 +201,7 @@ def jacobian(
       samples = samples.reshape(-1, samples.shape[-1])
       Or_k = jax.jacrev(lambda pars: logpsi(pars, samples).real, parameters)
       Oi_k = jax.jacrev(lambda pars: logpsi(pars, samples).imag, parameters)
-      O_k = jax.tree_map(lambda jr, ji: jnp.concatenate([jr, ji]], axis=1),
+      O_k = jax.tree_util.tree_map(lambda jr, ji: jnp.concatenate([jr, ji]], axis=1),
                                                         Or_k, Oi_k)
 
     As both :code:`Or_k` and :code:`Oi_k` are real, instead of concatenating we
@@ -220,7 +220,7 @@ def jacobian(
 
     .. code:: python
 
-      O_k_cmplx = jax.tree_map(lambda jri: jri[:, 0, :] + 1j* jri[:, 1, :], O_k)
+      O_k_cmplx = jax.tree_util.tree_map(lambda jri: jri[:, 0, :] + 1j* jri[:, 1, :], O_k)
 
 
     **If some parameters** :math:`\theta_k` **are complex**, this mode splits the
@@ -249,13 +249,13 @@ def jacobian(
 
       samples = samples.reshape(-1, samples.shape[-1])
       # tree_to_real splits the parameters in a tuple like
-      # {'real': jax.tree_map(jnp.real, pars), 'imag': jax.tree_map(jnp.imag, pars)}
+      # {'real': jax.tree.map(jnp.real, pars), 'imag': jax.tree.map(jnp.imag, pars)}
       pars_real, reconstruct = nk.jax.tree_to_real(parameters)
       Or_k = jax.jacrev(lambda pars_re: logpsi(reconstruct(pars_re), samples).real,
                         pars_real)
       Oi_k = jax.jacrev(lambda pars_re: logpsi(reconstruct(pars_re), samples).imag,
                         pars_real)
-      O_k = jax.tree_map(lambda jr, ji: jnp.concatenate([jr, ji]], axis=1),
+      O_k = jax.tree_util.tree_map(lambda jr, ji: jnp.concatenate([jr, ji]], axis=1),
                                                         Or_k, Oi_k)
 
     This code is also valid if all parameters are real, in which case :code:`O_k.real`
@@ -352,20 +352,24 @@ def jacobian(
 
     if pdf is None:
         if center:
-            jacobians = jax.tree_map(lambda x: subtract_mean(x, axis=0), jacobians)
+            jacobians = jax.tree_util.tree_map(
+                lambda x: subtract_mean(x, axis=0), jacobians
+            )
 
         if _sqrt_rescale:
             sqrt_n_samp = math.sqrt(
                 samples.shape[0] * mpi.n_nodes
             )  # maintain weak type
-            jacobians = jax.tree_map(lambda x: x / sqrt_n_samp, jacobians)
+            jacobians = jax.tree_util.tree_map(lambda x: x / sqrt_n_samp, jacobians)
 
     else:
         if center:
-            jacobians_avg = jax.tree_map(
+            jacobians_avg = jax.tree_util.tree_map(
                 partial(sum_mpi, axis=0), _multiply_by_pdf(jacobians, pdf)
             )
-            jacobians = jax.tree_map(lambda x, y: x - y, jacobians, jacobians_avg)
+            jacobians = jax.tree_util.tree_map(
+                lambda x, y: x - y, jacobians, jacobians_avg
+            )
 
         if _sqrt_rescale:
             jacobians = _multiply_by_pdf(jacobians, jnp.sqrt(pdf))
@@ -379,7 +383,7 @@ def _multiply_by_pdf(oks, pdf):
     Used to multiply the log-derivatives by the probability density.
     """
 
-    return jax.tree_map(
+    return jax.tree_util.tree_map(
         lambda x: jax.lax.broadcast_in_dim(pdf, x.shape, (0,)) * x,
         oks,
     )

{netket-3.11.3 → netket-3.11.4}/netket/jax/_scanmap.py

@@ -10,8 +10,10 @@ if jax.__version_info__ >= (0, 4, 16):
 else:
     from jax import linear_util as lu
 
-_tree_add = partial(jax.tree_map, jax.lax.add)
-_tree_zeros_like = partial(jax.tree_map, lambda x: jnp.zeros(x.shape, dtype=x.dtype))
+_tree_add = partial(jax.tree_util.tree_map, jax.lax.add)
+_tree_zeros_like = partial(
+    jax.tree_util.tree_map, lambda x: jnp.zeros(x.shape, dtype=x.dtype)
+)
 
 
 # TODO put it somewhere
@@ -62,7 +64,7 @@ def scan_append_reduce(f, x, append_cond, op=_tree_add, zero_fun=_tree_zeros_lik
     """
     # TODO: different op for each result
 
-    x0 = jax.tree_map(lambda x: x[0], x)
+    x0 = jax.tree_util.tree_map(lambda x: x[0], x)
 
     # special code path if there is only one element
     # to avoid having to rely on xla/llvm to optimize the overhead away

{netket-3.11.3 → netket-3.11.4}/netket/jax/_utils_random.py

@@ -36,7 +36,9 @@ def PRNGKey(
     else:
         key = seed
 
-    key = jax.tree_map(lambda k: mpi.mpi_bcast_jax(k, root=root, comm=comm)[0], key)
+    key = jax.tree_util.tree_map(
+        lambda k: mpi.mpi_bcast_jax(k, root=root, comm=comm)[0], key
+    )
 
     return key
 
@@ -90,7 +92,7 @@ def mpi_split(key, *, root=0, comm=MPI_jax_comm) -> PRNGKeyT:
     # on all MPI nodes?
     keys = jax.random.split(key, mpi.n_nodes)
 
-    keys = jax.tree_map(lambda k: mpi.mpi_bcast_jax(k, root=root)[0], keys)
+    keys = jax.tree_util.tree_map(lambda k: mpi.mpi_bcast_jax(k, root=root)[0], keys)
 
     return keys[mpi.rank]
 

{netket-3.11.3 → netket-3.11.4}/netket/jax/_utils_tree.py

@@ -76,14 +76,16 @@ def tree_leaf_iscomplex(pars: PyTree) -> bool:
     """
     Returns true if at least one leaf in the tree has complex dtype.
     """
-    return any(jax.tree_util.tree_leaves(jax.tree_map(jnp.iscomplexobj, pars)))
+    return any(
+        jax.tree_util.tree_leaves(jax.tree_util.tree_map(jnp.iscomplexobj, pars))
+    )
 
 
 def tree_leaf_isreal(pars: PyTree) -> bool:
     """
     Returns true if at least one leaf in the tree has real dtype.
     """
-    return any(jax.tree_util.tree_leaves(jax.tree_map(jnp.isrealobj, pars)))
+    return any(jax.tree_util.tree_leaves(jax.tree_util.tree_map(jnp.isrealobj, pars)))
 
 
 def tree_ishomogeneous(pars: PyTree) -> bool:
@@ -100,7 +102,9 @@ def tree_conj(t: PyTree) -> PyTree:
     Args:
         t: pytree
     """
-    return jax.tree_map(lambda x: jax.lax.conj(x) if jnp.iscomplexobj(x) else x, t)
+    return jax.tree_util.tree_map(
+        lambda x: jax.lax.conj(x) if jnp.iscomplexobj(x) else x, t
+    )
 
 
 @jax.jit
@@ -116,7 +120,9 @@ def tree_dot(a: PyTree, b: PyTree) -> Scalar:
     """
     return jax.tree_util.tree_reduce(
         jax.numpy.add,
-        jax.tree_map(jax.numpy.sum, jax.tree_map(jax.numpy.multiply, a, b)),
+        jax.tree_util.tree_map(
+            jax.numpy.sum, jax.tree_util.tree_map(jax.numpy.multiply, a, b)
+        ),
     )
 
 
@@ -135,7 +141,7 @@ def tree_cast(x: PyTree, target: PyTree) -> PyTree:
     """
     # astype alone would also work, however that raises ComplexWarning when casting complex to real
     # therefore the real is taken first where needed
-    return jax.tree_map(
+    return jax.tree_util.tree_map(
         lambda x, target: (x if jnp.iscomplexobj(target) else x.real).astype(
             target.dtype
         ),
@@ -157,9 +163,9 @@ def tree_axpy(a: Scalar, x: PyTree, y: PyTree) -> PyTree:
         where the leaves of x are first scaled with a.
     """
     if is_scalar(a):
-        return jax.tree_map(lambda x_, y_: a * x_ + y_, x, y)
+        return jax.tree_util.tree_map(lambda x_, y_: a * x_ + y_, x, y)
     else:
-        return jax.tree_map(lambda a_, x_, y_: a_ * x_ + y_, a, x, y)
+        return jax.tree_util.tree_map(lambda a_, x_, y_: a_ * x_ + y_, a, x, y)
 
 
 class RealImagTuple(tuple):
@@ -188,8 +194,8 @@ def _tree_to_real(x):
     if tree_leaf_iscomplex(x):
         # TODO find a way to make it a nop?
         # return jax.vmap(lambda y: jnp.array((y.real, y.imag)))(x)
-        r = jax.tree_map(lambda x: x.real if jnp.iscomplexobj(x) else x, x)
-        i = jax.tree_map(lambda x: x.imag if jnp.iscomplexobj(x) else None, x)
+        r = jax.tree_util.tree_map(lambda x: x.real if jnp.iscomplexobj(x) else x, x)
+        i = jax.tree_util.tree_map(lambda x: x.imag if jnp.iscomplexobj(x) else None, x)
         return RealImagTuple((r, i))
     else:
         return x
@@ -198,7 +204,9 @@ def _tree_to_real(x):
 def _tree_to_real_inverse(x):
     if isinstance(x, RealImagTuple):
         # not using jax.lax.complex because it would convert scalars to arrays
-        return jax.tree_map(lambda re, im: re + 1j * im if im is not None else re, *x)
+        return jax.tree_util.tree_map(
+            lambda re, im: re + 1j * im if im is not None else re, *x
+        )
     else:
         return x
 

{netket-3.11.3 → netket-3.11.4}/netket/jax/_vmap_chunked.py

@@ -20,12 +20,16 @@ def _eval_fun_in_chunks(vmapped_fun, chunk_size, argnums, *args, **kwargs):
     else:
         # split inputs
         def _get_chunks(x):
-            x_chunks = jax.tree_map(lambda x_: x_[: n_elements - n_rest, ...], x)
+            x_chunks = jax.tree_util.tree_map(
+                lambda x_: x_[: n_elements - n_rest, ...], x
+            )
             x_chunks = _chunk(x_chunks, chunk_size)
             return x_chunks
 
         def _get_rest(x):
-            x_rest = jax.tree_map(lambda x_: x_[n_elements - n_rest :, ...], x)
+            x_rest = jax.tree_util.tree_map(
+                lambda x_: x_[n_elements - n_rest :, ...], x
+            )
             return x_rest
 
         args_chunks = [
@@ -41,7 +45,9 @@ def _eval_fun_in_chunks(vmapped_fun, chunk_size, argnums, *args, **kwargs):
             y = y_chunks
         else:
             y_rest = vmapped_fun(*args_rest, **kwargs)
-            y = jax.tree_map(lambda y1, y2: jnp.concatenate((y1, y2)), y_chunks, y_rest)
+            y = jax.tree_util.tree_map(
+                lambda y1, y2: jnp.concatenate((y1, y2)), y_chunks, y_rest
+            )
     return y
 
 

{netket-3.11.3 → netket-3.11.4}/netket/jax/sharding.py

@@ -237,7 +237,7 @@ def extract_replicated(t):
         else:
             return x
 
-    return jax.tree_map(_extract_replicated, t)
+    return jax.tree_util.tree_map(_extract_replicated, t)
 
 
 def gather(x):
@@ -466,7 +466,9 @@ def sharding_decorator(f, sharded_args_tree, reduction_op_tree=False):
                     if o is True:
                         return lambda x: Partial(partial(lambda x: x, x))
                     else:
-                        return partial(jax.tree_map, partial(o, axis_name="i"))
+                        return partial(
+                            jax.tree_util.tree_map, partial(o, axis_name="i")
+                        )
 
                 reductions = [_sele_op(o) for o in reduction_op]
                 res = out_treedef.flatten_up_to(res)

{netket-3.11.3 → netket-3.11.4}/netket/optimizer/qgt/common.py

@@ -39,13 +39,15 @@ def check_valid_vector_type(x: PyTree, target: PyTree):
 
     try:
         if isinstance(target, RealImagTuple):
-            jax.tree_map(check, x, target.real, target.imag)
+            jax.tree_util.tree_map(check, x, target.real, target.imag)
         else:
-            jax.tree_map(check, x, target)
+            jax.tree_util.tree_map(check, x, target)
     except ValueError:
         # catches jax tree map errors
-        pars_struct = jax.tree_map(lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), x)
-        vec_struct = jax.tree_map(
+        pars_struct = jax.tree_util.tree_map(
+            lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), x
+        )
+        vec_struct = jax.tree_util.tree_map(
             lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), target
         )
 

{netket-3.11.3 → netket-3.11.4}/netket/optimizer/qgt/qgt_jacobian_common.py

@@ -82,7 +82,7 @@ def rescale(centered_oks, offset, *, ndims: int = 1):
     # should be (0,) for standard, (0,1) when we have 2 jacobians in complex mode
     axis = tuple(range(ndims))
 
-    scale = jax.tree_map(
+    scale = jax.tree_util.tree_map(
         lambda x: (
             mpi.mpi_sum_jax(jnp.sum((x * x.conj()).real, axis=axis, keepdims=True))[0]
             + offset
@@ -90,6 +90,6 @@ def rescale(centered_oks, offset, *, ndims: int = 1):
         ** 0.5,
         centered_oks,
     )
-    centered_oks = jax.tree_map(jnp.divide, centered_oks, scale)
-    scale = jax.tree_map(partial(jnp.squeeze, axis=axis), scale)
+    centered_oks = jax.tree_util.tree_map(jnp.divide, centered_oks, scale)
+    scale = jax.tree_util.tree_map(partial(jnp.squeeze, axis=axis), scale)
     return centered_oks, scale

{netket-3.11.3 → netket-3.11.4}/netket/optimizer/qgt/qgt_jacobian_dense.py

@@ -146,7 +146,7 @@ def QGTJacobianDense(
     else:
         scale = None
 
-    pars_struct = jax.tree_map(
+    pars_struct = jax.tree_util.tree_map(
         lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), vstate.parameters
     )
 

{netket-3.11.3 → netket-3.11.4}/netket/optimizer/qgt/qgt_jacobian_pytree.py

@@ -253,12 +253,12 @@ def _matmul(
         vec, reassemble = nkjax.tree_to_real(vec)
 
     if self.scale is not None:
-        vec = jax.tree_map(jnp.multiply, vec, self.scale)
+        vec = jax.tree_util.tree_map(jnp.multiply, vec, self.scale)
 
     result = mat_vec(vec, self.O, self.diag_shift)
 
     if self.scale is not None:
-        result = jax.tree_map(jnp.multiply, result, self.scale)
+        result = jax.tree_util.tree_map(jnp.multiply, result, self.scale)
 
     # Reassemble real-imaginary split as needed
     if reassemble is not None:
@@ -284,9 +284,9 @@ def _solve(
             x0, _ = nkjax.tree_to_real(x0)
 
     if self.scale is not None:
-        y = jax.tree_map(jnp.divide, y, self.scale)
+        y = jax.tree_util.tree_map(jnp.divide, y, self.scale)
         if x0 is not None:
-            x0 = jax.tree_map(jnp.multiply, x0, self.scale)
+            x0 = jax.tree_util.tree_map(jnp.multiply, x0, self.scale)
 
     # to pass the object LinearOperator itself down
     # but avoid rescaling, we pass down an object with
@@ -297,7 +297,7 @@ def _solve(
     out, info = solve_fun(unscaled_self, y, x0=x0)
 
     if self.scale is not None:
-        out = jax.tree_map(jnp.divide, out, self.scale)
+        out = jax.tree_util.tree_map(jnp.divide, out, self.scale)
 
     # Reassemble real-imaginary split as needed
     if self.mode != "holomorphic":
@@ -311,7 +311,7 @@ def _to_dense(self: QGTJacobianPyTreeT) -> jnp.ndarray:
     O = self.O
     if self.mode == "complex":
         # I want to iterate across the samples and real/imaginary part
-        O = jax.tree_map(lambda x: x.reshape(-1, *x.shape[2:]), O)
+        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)
 
     if self.scale is None:

{netket-3.11.3 → netket-3.11.4}/netket/optimizer/qgt/qgt_jacobian_pytree_logic.py

@@ -32,7 +32,7 @@ def _jvp(oks: PyTree, v: PyTree) -> Array:
     Compute the matrix-vector product between the pytree jacobian oks and the pytree vector v
     """
     td = lambda x, y: jnp.tensordot(x, y, axes=y.ndim)
-    t = jax.tree_map(td, oks, v)
+    t = jax.tree_util.tree_map(td, oks, v)
     return jax.tree_util.tree_reduce(jnp.add, t)
 
 
@@ -40,8 +40,8 @@ def _vjp(oks: PyTree, w: Array) -> PyTree:
     """
     Compute the vector-matrix product between the vector w and the pytree jacobian oks
     """
-    res = jax.tree_map(partial(jnp.tensordot, w, axes=w.ndim), oks)
-    return jax.tree_map(lambda x: mpi.mpi_sum_jax(x)[0], res)  # MPI
+    res = jax.tree_util.tree_map(partial(jnp.tensordot, w, axes=w.ndim), oks)
+    return jax.tree_util.tree_map(lambda x: mpi.mpi_sum_jax(x)[0], res)  # MPI
 
 
 def _mat_vec(v: PyTree, oks: PyTree) -> PyTree:

{netket-3.11.3 → netket-3.11.4}/netket/optimizer/qgt/qgt_onthefly.py

@@ -243,7 +243,7 @@ def _solve(
 
     # we could cache this...
     if x0 is None:
-        x0 = jax.tree_map(jnp.zeros_like, y)
+        x0 = jax.tree_util.tree_map(jnp.zeros_like, y)
 
     out, info = solve_fun(self, y, x0=x0)
     return out, info