NetKet 3.11.2__tar.gz → 3.11.4__tar.gz

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: NetKet
-Version: 3.11.2
+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/#
@@ -26,7 +26,7 @@ Requires-Dist: jaxlib<0.5,>=0.4.3
 Requires-Dist: numba4jax<0.1,>=0.0.10.post1
 Requires-Dist: numba<0.60,>=0.52
 Requires-Dist: numpy~=1.20
-Requires-Dist: optax<0.2,>=0.1.3
+Requires-Dist: optax<0.3,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
 Requires-Dist: rich>=10.0
 Requires-Dist: scipy<2,>=1.5.3
@@ -44,6 +44,7 @@ Requires-Dist: pytest-cov>=2.10.1; extra == 'dev'
 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.1.0; extra == 'dev'
 Requires-Dist: wheel; extra == 'dev'
 Provides-Extra: docs

{netket-3.11.2 → 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.2'
-__version_tuple__ = version_tuple = (3, 11, 2)
+__version__ = version = '3.11.4'
+__version_tuple__ = version_tuple = (3, 11, 4)

{netket-3.11.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → netket-3.11.4}/netket/experimental/observable/renyi2/exact.py

@@ -16,6 +16,7 @@ import numpy as np
 
 from netket.vqs import FullSumState, expect
 from netket.stats import Stats
+from netket.utils import module_version
 
 from .S2_operator import Renyi2EntanglementEntropy
 
@@ -37,6 +38,11 @@ def Renyi2(vstate: FullSumState, op: Renyi2EntanglementEntropy):
 
         rdm = state_qutip.ptrace(np.arange(N)[mask])
 
+        # TODO: when we drop support for qutip older versions, always keep this
+        # qutip 5 was released in march 2024
+        if module_version("qutip") >= (5, 0, 0):
+            rdm = rdm.data.to_array()
+
         n = 2
         out = np.log2(np.trace(np.linalg.matrix_power(rdm, n))) / (1 - n)
         out = np.absolute(out.real)

{netket-3.11.2 → netket-3.11.4}/netket/experimental/operator/_fermion_operator_2nd_base.py

@@ -19,9 +19,9 @@ import numpy as np
 from numbers import Number
 
 from netket.utils.types import DType
-from netket.operator._discrete_operator import DiscreteOperator
+from netket.operator import DiscreteOperator, Transpose
 from netket.operator._pauli_strings.base import _count_of_locations
-from netket.hilbert.abstract_hilbert import AbstractHilbert
+from netket.hilbert import AbstractHilbert
 from netket.utils.numbers import is_scalar, dtype as _dtype
 from netket.utils.optional_deps import import_optional_dependency
 
@@ -31,7 +31,7 @@ from ._fermion_operator_2nd_utils import (
     _convert_terms_to_spin_blocks,
     _canonicalize_input,
     _check_hermitian,
-    _herm_conj,
+    transpose_term,
     _remove_dict_zeros,
     _verify_input,
     OperatorDict,
@@ -267,10 +267,18 @@ class FermionOperator2ndBase(DiscreteOperator):
     def operators(self) -> OperatorDict:
         """Returns a dictionary with (term, weight) key-value pairs, with terms in tuple notation.
 
-        The constant diagonal shift of the operator is associated with a term represented as an empty tuple.
+        The constant diagonal shift of the operator is associated with a term represented as an
+        empty tuple.
         """
         return self._operators
 
+    @property
+    def cutoff(self) -> float:
+        """Threshold for the weights, if the absolute value of a weight is below the
+        cutoff, it's discarded.
+        """
+        return self._cutoff
+
     def copy(self, *, dtype: Optional[DType] = None, cutoff=None):
         """
         Creates a deep copy of this operator, potentially changing the dtype of the
@@ -458,23 +466,35 @@ class FermionOperator2ndBase(DiscreteOperator):
         op = self.copy(dtype=dtype)
         return op.__imul__(scalar)
 
+    def transpose(self, *, concrete=False):
+        r"""Returns the transpose of this operator.
+
+        For real operators, this is equivalent to complex conjugation.
+        This implementation always returns the concrete version of the operator.
+        """
+        if concrete:
+            new = type(self)(self.hilbert, dtype=self.dtype, cutoff=self._cutoff)
+            # operators (c,c†) are real already. Only conjugate coefficients if needed.
+            new._operators = {transpose_term(k): v for k, v in self._operators.items()}
+            return new
+        else:
+            return Transpose(self)
+
     def conjugate(self, *, concrete=False):
-        r"""Returns the complex conjugate of this operator."""
+        r"""Returns the complex conjugate of this operator.
 
-        terms = list(self._operators.keys())
-        weights = list(self._operators.values())
-        terms, weights = _herm_conj(terms, weights)  # changes also the terms
+        This implementation always returns the concrete version of the operator.
+        """
+        del concrete
 
-        cls = type(self)
-        new = cls(
-            self.hilbert,
-            dtype=self.dtype,
-        )
-        new._operators, _ = _canonicalize_input(
-            terms, weights, self.dtype, self._cutoff
-        )
-        new._cutoff = self._cutoff
-        return new
+        # if operator is real dtype, then just return a copy
+        if not np.issubdtype(self.dtype, np.complex_):
+            return self.copy()
+        else:
+            new = type(self)(self.hilbert, dtype=self.dtype, cutoff=self._cutoff)
+            # operators (c,c†) are real already. Only conjugate coefficients if needed.
+            new._operators = {k: np.conjugate(v) for k, v in self._operators.items()}
+            return new
 
     def to_normal_order(self):
         """Reoder the operators to normal order.

{netket-3.11.2 → netket-3.11.4}/netket/experimental/operator/_fermion_operator_2nd_utils.py

@@ -224,13 +224,29 @@ def _herm_conj(
     terms: OperatorTermsList, weights: OperatorWeightsList = 1
 ) -> tuple[OperatorTermsList, OperatorWeightsList]:
     """Returns the hermitian conjugate of the terms and weights."""
+    conj_term = transpose_terms(terms)
+    conj_weight = [np.conjugate(weight) for weight in weights]
+    return conj_term, conj_weight
+
+
+def transpose_terms(terms: OperatorTermsList) -> OperatorTermsList:
+    """Returns the transpose of the terms (creation/destruction operators).
+    The transpose is equivalent to the hermitian transpose because they are
+    real.
+    """
     conj_term = []
-    conj_weight = []
-    # loop over all terms and weights and get the hermitian conjugate
-    for term, weight in zip(terms, weights):
+    for term in terms:
         conj_term.append([(op, 1 - int(dag)) for (op, dag) in reversed(term)])
-        conj_weight.append(np.conjugate(weight))
-    return conj_term, conj_weight
+    return conj_term
+
+
+def transpose_term(term: OperatorTerm) -> OperatorTerm:
+    """Returns the transpose of a single term (creation/destruction operators).
+    The transpose is equivalent to the hermitian transpose because they are
+    real.
+    """
+    # equivalent to transpose_terms
+    return tuple((op, 1 - int(dag)) for (op, dag) in reversed(term))
 
 
 def _convert_terms_to_spin_blocks(

{netket-3.11.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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