quantax 0.2.0__tar.gz → 0.2.1__tar.gz

{quantax-0.2.0/quantax.egg-info → quantax-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quantax
-Version: 0.2.0
+Version: 0.2.1
 Summary: Flexible neural quantum states based on JAX
 Author: Christopher Roth
 Author-email: Ao Chen <chenao.phys@gmail.com>
@@ -8,7 +8,7 @@ Requires-Python: <=3.13,>=3.10
 License-File: LICENSE
 Requires-Dist: numpy>=1.26.0
 Requires-Dist: scipy>=1.11.1
-Requires-Dist: jax>=0.4.34
+Requires-Dist: jax>=0.6.1
 Requires-Dist: equinox>=0.11.4
 Requires-Dist: lrux>=0.1.0
 Provides-Extra: full
@@ -16,7 +16,7 @@ Requires-Dist: numpy>=2.0.0; extra == "full"
 Requires-Dist: scipy>=1.13.1; extra == "full"
 Requires-Dist: matplotlib>=3.8.0; extra == "full"
 Requires-Dist: quspin>=1.0.0; extra == "full"
-Requires-Dist: jax>=0.4.34; extra == "full"
+Requires-Dist: jax>=0.6.1; extra == "full"
 Requires-Dist: equinox>=0.11.4; extra == "full"
 Requires-Dist: lrux>=0.1.0; extra == "full"
 Dynamic: license-file

{quantax-0.2.0 → quantax-0.2.1}/README.md

@@ -1,34 +1,43 @@
 <h1 align='center'>Quantax</h1>
 
-**Flexible neural quantum states based on [JAX](https://github.com/google/jax)**
+<p align="center">
+  <strong>Flexible and scalable neural quantum states built on JAX</strong>
+</p>
 
-[📖 Documentation](https://chenao-phys.github.io/quantax)
+<p align="center">
+  <a href="https://chenao-phys.github.io/quantax">📖 Documentation</a>
+</p>
 
-## Not only NQS, but also
+---
 
-- Exact diagonalization (based on [QuSpin](https://github.com/QuSpin/QuSpin))
+
+## 🔍Scope
+
+Quantax is a research-oriented Python package for quantum many-body physics, with a focus on flexible and scalable neural quantum states (NQS). Apart from NQS, Quantax also includes several useful techniques in variational Monte Carlo (VMC):
+
+- Exact diagonalization and small-system benchmarks via [QuSpin](https://github.com/QuSpin/QuSpin)
 - Fermionic mean-field wavefunctions
-- Flexible neural network design (based on [Equinox](https://github.com/patrick-kidger/equinox))
+- Flexible (neural) wavefunction design via [Equinox](https://github.com/patrick-kidger/equinox)
+- Tensor networks via [quimb](https://github.com/jcmgray/quimb/tree/main) and [symmray](https://github.com/jcmgray/symmray) (to be implemented)
 
-## Installation
+## ⚙️Installation
 
-Requires Python 3.10+, JAX 0.4.34+
+Requires Python 3.10+, JAX 0.6.1+
 
-First, ensure that a correct JAX version is installed.
-For details, check [JAX Installation](https://docs.jax.dev/en/latest/installation.html).
+First, ensure that a correct JAX version is installed. For details, check [JAX Installation](https://docs.jax.dev/en/latest/installation.html).
 
 For a direct installation of full functionality (recommended in most cases),
-```
+```bash
 pip install quantax[full]
 ```
 
 For a minimal installation,
-```
+```bash
 pip install quantax
 ```
 
 
-## Quick Start
+## 🚀Quick Start
 
 ```python
 import quantax as qtx
@@ -65,3 +74,4 @@ for i in range(100):
 energy_data.plot(baseline=E)
 plt.show()
 ```
+![quick_start](images/quick_start.svg)

{quantax-0.2.0 → quantax-0.2.1}/docs/source/conf.py

@@ -14,7 +14,7 @@ sys.path.insert(0, pathlib.Path(__file__).parents[2].resolve().as_posix())
 project = "Quantax"
 copyright = "2025, Ao Chen"
 author = "Ao Chen, Christopher Roth"
-release = "0.2.0"
+release = "0.2.1"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

{quantax-0.2.0 → quantax-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quantax"
-version = "0.2.0"
+version = "0.2.1"
 description = "Flexible neural quantum states based on JAX"
 authors = [
     { name = "Ao Chen", email = "chenao.phys@gmail.com" },
@@ -14,7 +14,7 @@ requires-python = ">=3.10,<=3.13"
 dependencies = [
     "numpy>=1.26.0",
     "scipy>=1.11.1",
-    "jax>=0.4.34",
+    "jax>=0.6.1",
     "equinox>=0.11.4",
     "lrux>=0.1.0",
 ]
@@ -25,7 +25,7 @@ full = [
     "scipy>=1.13.1",
     "matplotlib>=3.8.0",
     "quspin>=1.0.0",
-    "jax>=0.4.34",
+    "jax>=0.6.1",
     "equinox>=0.11.4",
     "lrux>=0.1.0",
 ]

{quantax-0.2.0 → quantax-0.2.1}/quantax/model/conv_nets.py

@@ -93,6 +93,8 @@ class _ConvBlock(eqx.Module):
 
 
 class ResConv(Sequential):
+    """Deep convolutional residual network."""
+
     nblocks: int
     channels: int
     kernel_size: Union[int, Sequence[int]]
@@ -133,10 +135,15 @@ class ResConv(Sequential):
             The translation symmetry to be applied in the last layer, see `~quantax.nn.ConvSymmetrize`.
 
         :param dtype:
-            The data type of the parameters.
+            The data type of the parameters. Must be a real dtype.
+
+        :param out_dtype:
+            The data type of the output wavefunction. By default, it is the same as `dtype`.
+            If `out_dtype` is complex, `~quantax.nn.pair_cpl` will be applied to the output
+            of convolutional layers to make the final output complex.
 
         .. tip::
-            This is the recommended architecture for deep neural quantum states in spin systems.
+            This is the recommended architecture for deep NQS in spin systems.
         """
         if jnp.issubdtype(dtype, jnp.complexfloating):
             raise ValueError("`ResSum` doesn't support complex dtypes.")

{quantax-0.2.0 → quantax-0.2.1}/quantax/nn/modules.py

@@ -3,6 +3,7 @@ from typing import Sequence, Tuple, Union, Callable
 from jaxtyping import PyTree
 import jax
 import equinox as eqx
+from ..utils import PsiArray
 
 
 class Sequential(eqx.nn.Sequential):
@@ -35,7 +36,7 @@ class Sequential(eqx.nn.Sequential):
         super().__init__(layers)
         self.holomorphic = holomorphic
 
-    def __call__(self, x: jax.Array, *, s: jax.Array = None) -> jax.Array:
+    def __call__(self, x: jax.Array, *, s: jax.Array = None) -> PsiArray:
         """**Arguments:**
 
         - `x`: passed to the first member of the sequence.
@@ -74,7 +75,7 @@ class RawInputLayer(eqx.Module):
     basis state.
     """
 
-    def __call__(self, x: jax.Array, s: jax.Array) -> jax.Array:
+    def __call__(self, x: jax.Array, s: jax.Array) -> Callable:
         """
         The forward pass.
 
@@ -97,7 +98,7 @@ class RefModel(eqx.Module):
         Return initial internal values for the given configuration.
         """
 
-    def __call__(self, x: jax.Array) -> jax.Array:
+    def __call__(self, x: jax.Array) -> PsiArray:
         """
         Usual forward pass without internal quantities.
         """
@@ -109,7 +110,7 @@ class RefModel(eqx.Module):
         nflips: int,
         internal: PyTree,
         return_update: bool = False,
-    ) -> Union[jax.Array, Tuple[jax.Array, PyTree]]:
+    ) -> Union[PsiArray, Tuple[PsiArray, PyTree]]:
         """
         Accelerated forward pass through local updates and internal quantities.
 

{quantax-0.2.0 → quantax-0.2.1}/quantax/optimizer/__init__.py

@@ -15,4 +15,4 @@ from .solver import (
     minsr_pinv_eig,
     sgd_solver,
 )
-from .time_evol import TimeEvol, ExactTimeEvol
+from .time_evol import TimeEvol

{quantax-0.2.0 → quantax-0.2.1}/quantax/optimizer/driver.py

@@ -7,7 +7,7 @@ from ..sampler import Sampler
 from ..utils import DataTracer
 
 
-class Driver():
+class Driver:
     def __init__(
         self,
         state: Variational,
@@ -26,8 +26,10 @@ class Driver():
     def step(self) -> None:
         """iteration"""
 
+
 class Euler(Driver):
     """First order Euler driver"""
+
     def step(self) -> None:
         samples = self._sampler.sweep()
         step = self._optimizer.get_step(samples)
@@ -39,6 +41,7 @@ class Euler(Driver):
 
 class AdaptiveHeunEvolution(Driver):
     """Adaptive second order Heun driver, designed for unitary time-evolution"""
+
     def __init__(
         self,
         state: Variational,
@@ -57,13 +60,13 @@ class AdaptiveHeunEvolution(Driver):
 
         samples = self._sampler.sweep()
         stepi = tdvp.get_step(samples)
-        self._state.update(dt * stepi, rescale=False)
+        self._state.update(dt * stepi)
         dtheta0_i = 1j * tdvp.energy - jnp.dot(tdvp._Omean, stepi)
-        
+
         samples = self._sampler.sweep()
         stepf = tdvp.get_step(samples)
         step1 = (stepi + stepf) / 2
-        self._state.update(-dt / 2 * stepi, rescale=False)
+        self._state.update(-dt / 2 * stepi)
 
         samples = self._sampler.sweep()
         stepm = tdvp.get_step(samples)
@@ -72,18 +75,18 @@ class AdaptiveHeunEvolution(Driver):
 
         samples = self._sampler.sweep()
         stepmm = tdvp.get_step(samples)
-        self._state.update(dt / 2 * stepmm, rescale=False)
+        self._state.update(dt / 2 * stepmm)
         dtheta0_mm = 1j * tdvp.energy - jnp.dot(tdvp._Omean, stepmm)
 
         samples = self._sampler.sweep()
         Smat, Fvec = tdvp.get_SF(samples)
         stepff = tdvp.solve(Smat, Fvec)
-        self._state.update(dt / 4 * (stepff - stepmm), rescale=True)
+        self._state.update(dt / 4 * (stepff - stepmm))
         dtheta0_ff = 1j * tdvp.energy - jnp.dot(tdvp._Omean, stepff)
         step2 = (stepi + stepm + stepmm + stepff) / 4
 
-        dtheta0 = dt * (dtheta0_i + dtheta0_m + dtheta0_mm + dtheta0_ff) / 4
-        self._state.rescale(dtheta0)
+        # dtheta0 = dt * (dtheta0_i + dtheta0_m + dtheta0_mm + dtheta0_ff) / 4
+        # self._state.rescale(dtheta0)
 
         diff = step1 - step2
         new_err = jnp.einsum("k,kq,q", diff.conj(), Smat, diff).real
@@ -95,5 +98,5 @@ class AdaptiveHeunEvolution(Driver):
         self._step_length = new_step_length
 
         self.step_size.append(new_step_length, self._time)
-        self.energy.append(self._optimizer.energy, self._time)
-        self.VarE.append(self._optimizer.VarE, self._time)
+        self.energy.append(tdvp.energy, self._time)
+        self.VarE.append(tdvp.VarE, self._time)

{quantax-0.2.0 → quantax-0.2.1}/quantax/optimizer/sr.py

@@ -491,7 +491,7 @@ class ER(QNGD):
         energy = dense @ H_psi
         Ebar = H_psi - dense * energy
         self._energy = energy.real
-        return Ebar.psi
+        return jnp.asarray(Ebar.psi)
 
     def get_Obar(self, psi: jax.Array) -> jax.Array:
         r"""Compute :math:`\bar O` in the full Hilbert space."""

{quantax-0.2.0 → quantax-0.2.1}/quantax/optimizer/time_evol.py

@@ -1,13 +1,11 @@
-from typing import Union, Tuple, Optional, Callable
-from numbers import Number
+from typing import Tuple, Optional, Callable
 from functools import partial
-import numpy as np
 import jax
 import jax.numpy as jnp
 
 from .sr import SR
 from .solver import pinvh_solve
-from ..state import DenseState, Variational, VS_TYPE
+from ..state import Variational, VS_TYPE
 from ..operator import Operator
 from ..sampler import Samples
 from ..utils import get_distribute_sharding, array_extend
@@ -26,18 +24,40 @@ def _AconjB(A: jax.Array, B: jax.Array) -> jax.Array:
 
 
 class TimeEvol(SR):
+    r"""
+    Time evolution optimizer, equivalent to real-time `~quantax.optimizer.SR`.
+    This optimizer assumes the number of samples is more than the number of parameters,
+    and is more memory-efficient than `~quantax.optimizer.SR` when this is the case.
+    """
+
     def __init__(
         self,
         state: Variational,
         hamiltonian: Operator,
         solver: Optional[Callable] = None,
     ):
+        r"""
+        :param state:
+            Variational state to be evolved.
+
+        :param hamiltonian:
+            Hamiltonian operator for time evolution.
+
+        :param solver:
+            The numerical solver for :math:`Sx = F`, default to pseudo-inverse.
+        """
+
         if solver is None:
             solver = pinvh_solve()
         super().__init__(state, hamiltonian, imag_time=False, solver=solver)
         self._max_parallel = state._backward_chunk
 
     def get_SF(self, samples: Samples) -> Tuple[jax.Array, jax.Array]:
+        r"""
+        Compute :math:`S = \bar O^\dagger \bar O` and :math:`F = \bar O^\dagger \bar \epsilon`
+        with the given samples. When the number of samples is large, this function will
+        automatically switch to a more memory-efficient implementation.
+        """
         if (
             self._max_parallel is None
             or samples.nsamples <= self._max_parallel * jax.device_count()
@@ -81,7 +101,6 @@ class TimeEvol(SR):
             newF = _AconjB(Omat, e)
             Smat += newS
             Fvec += newF
-        # psum here, nsamples definition should be modified to all samples across nodes
         Smat = jnp.sum(Smat, axis=0) / nsamples
         Fvec = jnp.sum(Fvec, axis=0) / nsamples
         Omean = jnp.sum(Omean, axis=0) / nsamples
@@ -113,36 +132,3 @@ class TimeEvol(SR):
         Smat, Fvec = self.get_SF(samples)
         step = self.solve(Smat, Fvec)
         return step
-
-
-class ExactTimeEvol:
-    def __init__(self, init_state: DenseState, hamiltonian: Operator) -> None:
-        self._init_state = init_state
-        self._symm = self._init_state.symm
-        self._eigs, self._U = hamiltonian.diagonalize(self._symm, "full")
-        self._eigs = jnp.asarray(self._eigs)
-        self._U = jnp.asarray(self._U)
-
-    def get_evolved_psi(self, time: Union[float, jax.Array]) -> jax.Array:
-        is_float = isinstance(time, float)
-        if is_float:
-            time = jnp.array([time])
-        exp_eigs = jnp.exp(-1j * jnp.einsum("t,d->td", time, self._eigs))
-        psi0 = self._init_state.psi
-        psi = jnp.einsum("ij,tj,kj,k->ti", self._U, exp_eigs, self._U.conj(), psi0)
-        if is_float:
-            psi = psi[0]
-        return psi
-
-    def expectation(
-        self, operator: Operator, time: Union[float, jax.Array]
-    ) -> Union[Number, jax.Array]:
-        psi = self.get_evolved_psi(time)
-        psi /= jnp.linalg.norm(psi, axis=1, keepdims=True)
-        op = operator.get_quspin_op(self._symm)
-
-        # this hasn't been tested
-        out = jnp.einsum("ti,it->t", psi.conj(), op.dot(np.ascontiguousarray(psi.T)))
-        if isinstance(time, float):
-            out = out.item()
-        return out

{quantax-0.2.0 → quantax-0.2.1}/quantax/sampler/metropolis.py

@@ -181,15 +181,10 @@ class Metropolis(Sampler):
 
         state = self._state
         use_ref = state.use_ref and self.nflips is not None
-        if use_ref:
-            has_chunk = hasattr(state, "ref_chunk") and state.ref_chunk is not None
-            chunk_size = state.ref_chunk
-        else:
-            has_chunk = (
-                hasattr(state, "forward_chunk") and state.forward_chunk is not None
-            )
-            chunk_size = state.forward_chunk
-        is_chunked = has_chunk and chunk_size < self.nsamples // jax.device_count()
+        attr = "ref_chunk" if use_ref else "forward_chunk"
+        chunk_size = getattr(state, attr, None)
+        ns = self.nsamples // jax.device_count()
+        is_chunked = chunk_size is not None and chunk_size < ns
 
         if is_chunked:
             if use_ref:

{quantax-0.2.0 → quantax-0.2.1}/quantax/state/variational.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Callable, Optional, Tuple, Union, BinaryIO
+from typing import Optional, Tuple, Union, BinaryIO
 from jaxtyping import PyTree
 from pathlib import Path
 
@@ -16,6 +16,8 @@ from .state import State
 from ..symmetry import Symmetry
 from ..nn import RefModel
 from ..utils import (
+    chunk_map,
+    shard_vmap,
     chunk_shard_vmap,
     to_distribute_array,
     filter_replicate,
@@ -114,7 +116,7 @@ class Variational(State):
         model: eqx.Module,
         param_file: Optional[Union[str, Path, BinaryIO]] = None,
         symm: Optional[Symmetry] = None,
-        max_parallel: Union[None, int, Tuple[int, int]] = None,
+        max_parallel: Union[None, int, Tuple[int, int], Tuple[int, int, int]] = None,
         use_ref: bool = True,
     ):
         r"""
@@ -126,23 +128,35 @@ class Variational(State):
             or `equinox.tree_serialise_leaves`, default to not loading parameters.
 
         :param symm: Symmetry of the network, default to `~quantax.symmetry.Identity`.
+            Denoting the network output as :math:`f(s)`
+            and symmetry elements as :math:`T_i` with characters :math:`\omega_i`,
+            the wave function is given by
+            :math:`\psi(s) = \sum_i \omega_i \, f(T_i s) / n_{symm}`
 
         :param max_parallel:
-            The maximum foward pass allowed per device, default to no limit.
-            Specifying a limited value is important for large batches to avoid memory overflow.
-            For Heisenberg-like hamiltonian, this also helps to improve the efficiency
-            of computing local energy by keeping constant amount of forward pass and
-            avoiding re-jitting.
+            The maximum chunk size allowed per device.
+            Specifying a limited value is important for avoiding memory overflow.
+            For many hamiltonians, this also helps to improve the efficiency
+            by keeping a constant amount of forward pass and avoiding re-jitting.
+            The allowed input formats are:
+
+            - None:
+                No chunk size (default).
+
+            - int:
+                The same chunk size for all forward and backward passes.
+
+            - Tuple[int, int]:
+                The chunk size for forward and backward passes respectively.
+
+            - Tuple[int, int, int]:
+                The chunk size for forward pass, backward pass and
+                `~quantax.state.Variational.ref_forward_with_updates` respectively.
 
         :param use_ref:
             Whether `ref_forward` and `ref_forward_with_updates` will be used when
             the model is a `~quantax.nn.RefModel`. When the model is not a `RefModel`,
             this argument has no effect. Default to ``True``.
-
-        Denoting the network output as :math:`f(s)`,
-        and symmetry elements as :math:`T_i` with characters :math:`\omega_i`,
-        the final wave function is given by
-        :math:`\psi(s) = \sum_i \omega_i \, f(T_i s) / n_{symm}`
         """
         super().__init__(symm)
         if param_file is not None:
@@ -180,17 +194,20 @@ class Variational(State):
 
     @property
     def forward_chunk(self) -> int:
-        """The maximum foward pass allowed per device."""
+        """The maximum chunk size of forward pass allowed per device."""
         return self._forward_chunk
 
     @property
     def backward_chunk(self) -> int:
-        """The maximum backward pass allowed per device."""
+        """The maximum chunk size of backward pass allowed per device."""
         return self._backward_chunk
 
     @property
     def ref_chunk(self) -> int:
-        """The maximum reference forward with updates allowed per device."""
+        """
+        The maximum chunk size of `~quantax.state.Variational.ref_forward_with_updates`
+        allowed per device.
+        """
         return self._ref_chunk
 
     @property
@@ -242,23 +259,28 @@ class Variational(State):
             )
 
     def _init_forward(self) -> None:
-        def direct_forward(model: eqx.Module, s: jax.Array) -> jax.Array:
+        def batch_forward(model: eqx.Module, s: jax.Array) -> jax.Array:
             s_symm = self.symm.get_symm_spins(s)
             psi = jax.vmap(model)(s_symm)
             psi = self.symm.symmetrize(psi, s)
             return psi.astype(get_default_dtype())
 
-        self._direct_forward = chunk_shard_vmap(
-            direct_forward, in_axes=(None, 0), out_axes=0, chunk_size=self.forward_chunk
+        self._batch_forward = shard_vmap(batch_forward, in_axes=(None, 0), out_axes=0)
+        self._direct_forward = chunk_map(
+            self._batch_forward, in_axes=(None, 0), chunk_size=self.forward_chunk
+        )
+        self._fulljit_forward = chunk_shard_vmap(
+            batch_forward, in_axes=(None, 0), out_axes=0, chunk_size=self.forward_chunk
         )
 
         def init_internal(model, s):
             s_symm = self.symm.get_symm_spins(s)
             return jax.vmap(model.init_internal)(s_symm)
 
-        self._init_internal = chunk_shard_vmap(
+        init_internal = chunk_shard_vmap(
             init_internal, in_axes=(None, 0), out_axes=0, chunk_size=self.ref_chunk
         )
+        self._init_internal = eqx.filter_jit(init_internal)
 
         def ref_forward_with_updates(model, s, s_old, nflips, internal):
             s_symm = self.symm.get_symm_spins(s)
@@ -288,11 +310,15 @@ class Variational(State):
             psi = self.symm.symmetrize(psi, s)
             return psi.astype(get_default_dtype())
 
-        self._ref_forward = chunk_shard_vmap(
+        self._batch_ref_forward = shard_vmap(
             ref_forward,
             in_axes=(None, 0, None, None, 0, None),
-            shard_axes=(None, 0, 0, None, 0, 0),
             out_axes=0,
+            shard_axes=(None, 0, 0, None, 0, 0),
+        )
+        self._ref_forward = chunk_map(
+            self._batch_ref_forward,
+            in_axes=(None, 0, None, None, 0, None),
             chunk_size=self.forward_chunk,
         )
 
@@ -355,7 +381,7 @@ class Variational(State):
         if self._use_ref:
             out = self._ref_forward_with_updates(self.model, s, s_old, nflips, internal)
         else:
-            out = self(s), None
+            out = self._fulljit_forward(self.model, s), None
         return out
 
     def ref_forward(
@@ -392,7 +418,7 @@ class Variational(State):
         if self._use_ref:
             out = self._ref_forward(self.model, s, s_old, nflips, idx_segment, internal)
         else:
-            out = self(s)
+            out = self._direct_forward(self.model, s)
         return out
 
     def _init_backward(self) -> None:

{quantax-0.2.0 → quantax-0.2.1}/quantax/utils/__init__.py

@@ -23,5 +23,5 @@ from .tree import (
     apply_updates,
 )
 from .big_array import LogArray, ScaleArray, PsiArray, where
-from .function import chunk_shard_vmap, chunk_map, shmap
+from .function import chunk_map, shmap, shard_vmap, chunk_shard_vmap
 from .basis import ints_to_array, array_to_ints, neel, stripe, Sqz_factor, rand_states

{quantax-0.2.0 → quantax-0.2.1}/quantax/utils/function.py

@@ -4,7 +4,6 @@ import jax
 import jax.numpy as jnp
 from jaxtyping import PyTree
 from jax.sharding import Mesh, PartitionSpec
-from jax.experimental.shard_map import shard_map
 import equinox as eqx
 from .array import array_extend
 from .tree import filter_tree_map
@@ -180,19 +179,20 @@ def shmap(
     mesh = Mesh(jax.devices(), "x")
     in_specs = _axes_to_specs(in_axes)
     out_specs = _axes_to_specs(out_axes)
-    f = shard_map(f, mesh, in_specs, out_specs, check_rep=False)
+    f = jax.shard_map(
+        f, mesh=mesh, in_specs=in_specs, out_specs=out_specs, check_vma=False
+    )
     return f
 
 
-def chunk_shard_vmap(
+def shard_vmap(
     f: Callable,
     in_axes: Union[tuple, int, None],
     out_axes: Union[tuple, int, None],
-    chunk_size: Optional[int] = None,
     shard_axes: Union[tuple, int, None] = None,
 ) -> Callable:
     """
-    f -> jit(chunk_map(shard_map(vmap(f))))
+    f -> shard_map(vmap(f))
 
     :param f:
         The function to be converted. The arguments of f will be sharded.
@@ -202,12 +202,44 @@ def chunk_shard_vmap(
 
     :param out_axes:
         The mapped axes of f outputs.
-    """
 
+    :param shard_axes:
+        The sharded axes of f input arguments. Default to in_axes.
+    """
     shard_axes = in_axes if shard_axes is None else shard_axes
 
     f = jax.vmap(f, in_axes, out_axes)
     f = shmap(f, shard_axes, out_axes)
+    f = eqx.filter_jit(f)
+    return f
+
+
+def chunk_shard_vmap(
+    f: Callable,
+    in_axes: Union[tuple, int, None],
+    out_axes: Union[tuple, int, None],
+    shard_axes: Union[tuple, int, None] = None,
+    chunk_size: Optional[int] = None,
+) -> Callable:
+    """
+    f -> chunk_map(shard_map(vmap(f)), use_scan=True)
+
+    :param f:
+        The function to be converted. The arguments of f will be sharded.
+
+    :param in_axes:
+        The mapped axes of f input arguments.
+
+    :param out_axes:
+        The mapped axes of f outputs.
+
+    :param shard_axes:
+        The sharded axes of f input arguments. Default to in_axes.
+
+    :param chunk_size:
+        The chunk size on each machine. If None, no chunking will be applied.
+    """
+    f = shard_vmap(f, in_axes, out_axes, shard_axes)
     f = chunk_map(f, in_axes, out_axes, chunk_size, use_scan=True)
     f = eqx.filter_jit(f)
     return f

{quantax-0.2.0 → quantax-0.2.1/quantax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quantax
-Version: 0.2.0
+Version: 0.2.1
 Summary: Flexible neural quantum states based on JAX
 Author: Christopher Roth
 Author-email: Ao Chen <chenao.phys@gmail.com>
@@ -8,7 +8,7 @@ Requires-Python: <=3.13,>=3.10
 License-File: LICENSE
 Requires-Dist: numpy>=1.26.0
 Requires-Dist: scipy>=1.11.1
-Requires-Dist: jax>=0.4.34
+Requires-Dist: jax>=0.6.1
 Requires-Dist: equinox>=0.11.4
 Requires-Dist: lrux>=0.1.0
 Provides-Extra: full
@@ -16,7 +16,7 @@ Requires-Dist: numpy>=2.0.0; extra == "full"
 Requires-Dist: scipy>=1.13.1; extra == "full"
 Requires-Dist: matplotlib>=3.8.0; extra == "full"
 Requires-Dist: quspin>=1.0.0; extra == "full"
-Requires-Dist: jax>=0.4.34; extra == "full"
+Requires-Dist: jax>=0.6.1; extra == "full"
 Requires-Dist: equinox>=0.11.4; extra == "full"
 Requires-Dist: lrux>=0.1.0; extra == "full"
 Dynamic: license-file

{quantax-0.2.0 → quantax-0.2.1}/quantax.egg-info/requires.txt

@@ -1,6 +1,6 @@
 numpy>=1.26.0
 scipy>=1.11.1
-jax>=0.4.34
+jax>=0.6.1
 equinox>=0.11.4
 lrux>=0.1.0
 
@@ -9,6 +9,6 @@ numpy>=2.0.0
 scipy>=1.13.1
 matplotlib>=3.8.0
 quspin>=1.0.0
-jax>=0.4.34
+jax>=0.6.1
 equinox>=0.11.4
 lrux>=0.1.0