NetKet 3.14.3__tar.gz → 3.14.4__tar.gz

{netket-3.14.3 → netket-3.14.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.14.3
+Version: 3.14.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/#
@@ -18,24 +18,24 @@ Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.10
-Requires-Dist: equinox<0.12,>=0.11.3
-Requires-Dist: flax<0.10,>=0.7.1
-Requires-Dist: igraph<0.12.0,>=0.10.0
-Requires-Dist: jax<0.5,>=0.4.16
-Requires-Dist: numba<0.61,>=0.52
+Requires-Dist: equinox>=0.11.3
+Requires-Dist: flax>=0.7.1
+Requires-Dist: igraph>=0.10.0
+Requires-Dist: jax>=0.4.16
+Requires-Dist: numba>=0.52
 Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
-Requires-Dist: plum-dispatch<3,>=2.4
+Requires-Dist: plum-dispatch>=2.4
 Requires-Dist: rich>=10.0
 Requires-Dist: scipy<2,>=1.5.3
 Requires-Dist: tqdm<5,>=4.60
 Provides-Extra: dev
 Requires-Dist: black==23.9.1; extra == 'dev'
-Requires-Dist: build; extra == 'dev'
+Requires-Dist: build>=1; extra == 'dev'
 Requires-Dist: coverage>=5; extra == 'dev'
 Requires-Dist: matplotlib>=3; extra == 'dev'
-Requires-Dist: networkx<4,>=2.4; extra == 'dev'
+Requires-Dist: networkx>=2.4; extra == 'dev'
 Requires-Dist: pre-commit>=2.7; extra == 'dev'
 Requires-Dist: pytest-cov>=2.10.1; extra == 'dev'
 Requires-Dist: pytest-json-report>=1.3; extra == 'dev'
@@ -43,14 +43,14 @@ Requires-Dist: pytest-xdist[psutil]>=2; extra == 'dev'
 Requires-Dist: pytest>=6; extra == 'dev'
 Requires-Dist: qutip>=5; extra == 'dev'
 Requires-Dist: ruff==0.6.2; extra == 'dev'
-Requires-Dist: wheel; extra == 'dev'
+Requires-Dist: wheel>=0.38.1; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: myst-nb<1.2,>=1.0; extra == 'docs'
 Requires-Dist: myst-parser<4.1.0,>=2.0.0; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints<3.0,>=1.22; extra == 'docs'
 Requires-Dist: sphinx-book-theme~=1.1.0; extra == 'docs'
 Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
-Requires-Dist: sphinx<8.1,>=7.0; extra == 'docs'
+Requires-Dist: sphinx<8.2,>=7.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
 Provides-Extra: extra
@@ -59,7 +59,7 @@ Requires-Dist: openfermion>=1.0.0; extra == 'extra'
 Requires-Dist: qutip>=4; extra == 'extra'
 Requires-Dist: tensorboardx>=2.0.0; extra == 'extra'
 Provides-Extra: mpi
-Requires-Dist: mpi4jax<0.6,>=0.3.9; extra == 'mpi'
+Requires-Dist: mpi4jax>=0.3.9; extra == 'mpi'
 Requires-Dist: mpi4py<5,>=3.0.1; extra == 'mpi'
 Provides-Extra: pyscf
 Requires-Dist: pyscf>=2.0; extra == 'pyscf'

{netket-3.14.3 → netket-3.14.4}/netket/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.14.3'
-__version_tuple__ = version_tuple = (3, 14, 3)
+__version__ = version = '3.14.4'
+__version_tuple__ = version_tuple = (3, 14, 4)

{netket-3.14.3 → netket-3.14.4}/netket/hilbert/constraint/base.py

@@ -100,7 +100,7 @@ class DiscreteHilbertConstraint(struct.Pytree):
                     return jax.pure_callback(self._call_py,
                                             result_shape_dtypes=(jax.ShapeDtypeStruct(x.shape[:-1], bool)),
                                             x,
-                                            vectorized=True)
+                                            vmap_method="expand_dims")
 
                 def _call_py(self, x):
                     # Not Jax compatible

{netket-3.14.3 → netket-3.14.4}/netket/hilbert/index/constrained_sum.py

@@ -78,21 +78,21 @@ class SumConstrainedHilbertIndex(HilbertIndex):
     @jax.jit
     def _compute_all_states(self):
         if self.n_particles == 0:
-            return jnp.zeros((1, self.size), dtype=jnp.int32)
+            return jnp.zeros((1, self.size), dtype=self.range.dtype)
         with jax.ensure_compile_time_eval():
             c = jnp.repeat(
-                jnp.eye(self.size, dtype=jnp.int32),
+                jnp.eye(self.size, dtype=self.range.dtype),
                 np.array(self.shape) - 1,
                 axis=0,
             )
             combs = jnp.array(
                 list(itertools.combinations(np.arange(len(c)), self.n_particles))
             )
-            all_states = c[combs].sum(axis=1, dtype=jnp.int32)
+            all_states = c[combs].sum(axis=1, dtype=self.range.dtype)
             if (np.array(self.shape) > 1).any():
                 all_states = jnp.unique(all_states, axis=0)
         all_states_fock = jnp.asarray(all_states)
-        return self.range.numbers_to_states(all_states_fock, dtype=np.int32)
+        return self.range.numbers_to_states(all_states_fock, dtype=self.range.dtype)
 
     @struct.property_cached(pytree_node=True)
     def _lookup_table(self) -> LookupTableHilbertIndex:

{netket-3.14.3 → netket-3.14.4}/netket/sampler/rules/hamiltonian.py

@@ -123,7 +123,7 @@ class HamiltonianRuleNumba(HamiltonianRuleBase):
             ),
             σ,
             rand_vec,
-            # vectorized=True, #TODO: Make the callback work with vectorized inputs
+            # vmap_method="expand_dims", #TODO: Make the callback work with vectorized inputs
         )
 
         return σp, log_prob_correction

{netket-3.14.3 → netket-3.14.4}/netket/vqs/mc/kernels.py

@@ -59,6 +59,7 @@ def local_value_kernel(logpsi: Callable, pars: PyTree, σ: Array, args: PyTree):
     σp, mel = args
     return jnp.sum(mel * jnp.exp(logpsi(pars, σp) - logpsi(pars, σ)))
 
+
 def local_value_kernel_jax(
     logpsi: Callable, pars: PyTree, σ: Array, O: DiscreteJaxOperator
 ):
@@ -70,21 +71,28 @@ def local_value_kernel_jax(
     logpsi_σp = logpsi(pars, σp.reshape(-1, σp.shape[-1])).reshape(σp.shape[:-1])
     return jnp.sum(mel * jnp.exp(logpsi_σp - jnp.expand_dims(logpsi_σ, -1)), axis=-1)
 
+
 def local_value_kernel_jax_conn_chunked(
-    logpsi: Callable, pars: PyTree, σ: Array, O: DiscreteJaxOperator, chunk_size: int,
+    logpsi: Callable,
+    pars: PyTree,
+    σ: Array,
+    O: DiscreteJaxOperator,
+    chunk_size: int,
 ):
     """
     local_value kernel for MCState for jax-compatible operators
     """
+    apply_conn = lambda s: logpsi(pars, s)
+    apply_conn = nkjax.apply_chunked(apply_conn, in_axes=0, chunk_size=chunk_size)
 
-    σp, mel = O.get_conn_padded(jnp.squeeze(σ))
-    logpsi_σ = logpsi(pars, σ)
-    apply_conn = lambda s: logpsi(pars,s)
-    apply_conn = nkjax.apply_chunked(apply_conn,in_axes=0,chunk_size=chunk_size) 
-    logpsi_σp = apply_conn(σp)
+    σp, mel = O.get_conn_padded(σ)
+
+    logpsi_σ = apply_conn(σ)
+    logpsi_σp = apply_conn(σp.reshape(-1, σ.shape[-1])).reshape(σp.shape[:-1])
 
     return jnp.sum(mel * jnp.exp(logpsi_σp - jnp.expand_dims(logpsi_σ, -1)), axis=-1)
 
+
 def local_value_squared_kernel(logpsi: Callable, pars: PyTree, σ: Array, args: PyTree):
     """
     local_value kernel for MCState and Squared (generic) operators
@@ -196,15 +204,16 @@ def local_value_kernel_jax_chunked(
     """
     local_value kernel for MCState and jaxcoompatible operators
     """
-
     if chunk_size >= O.max_conn_size:
         local_value_kernel = lambda s: local_value_kernel_jax(logpsi, pars, s, O)
+        local_value_chunked = nkjax.apply_chunked(
+            local_value_kernel,
+            in_axes=0,
+            chunk_size=max(1, chunk_size // O.max_conn_size),
+        )
     else:
-        local_value_kernel = lambda s: local_value_kernel_jax_conn_chunked(logpsi, pars, s, O, chunk_size)
-    
-    local_value_chunked = nkjax.apply_chunked(
-        local_value_kernel, in_axes=0, chunk_size=max(1, chunk_size // O.max_conn_size)
-    )
+        local_value_chunked = lambda s: local_value_kernel_jax_conn_chunked(
+            logpsi, pars, s, O, chunk_size
+        )
 
     return local_value_chunked(σ)
-

{netket-3.14.3 → netket-3.14.4}/pyproject.toml

@@ -35,20 +35,20 @@ dependencies = [
     "scipy>=1.5.3, <2",
     "tqdm>=4.60, <5",
     "rich>=10.0",
-    "numba>=0.52, <0.61",
-    "igraph>=0.10.0, <0.12.0",
-    "jax>=0.4.16, <0.5",
-    "equinox>=0.11.3, <0.12",
-    "flax>=0.7.1, <0.10",
+    "numba>=0.52",
+    "igraph>=0.10.0",
+    "jax>=0.4.16",
+    "equinox>=0.11.3",
+    "flax>=0.7.1",
     "orjson>=3.4, <4",
     "optax>=0.1.3, <0.3",
-    "plum-dispatch>=2.4, <3",
+    "plum-dispatch>=2.4",
 ]
 
 [project.optional-dependencies]
 mpi = [
     "mpi4py>=3.0.1, <5",
-    "mpi4jax>=0.3.9, <0.6"
+    "mpi4jax>=0.3.9"
     ]
 pyscf = [
     "sparse>=0.12.0",
@@ -61,7 +61,7 @@ extra = [
     "qutip>=4",
     ]
 dev = [
-    "networkx>=2.4, <4",
+    "networkx>=2.4",
     "matplotlib>=3",
     "pytest>=6",
     "pytest-xdist[psutil]>=2",
@@ -72,11 +72,11 @@ dev = [
     "black==23.9.1",
     "ruff==0.6.2",
     "qutip>=5",
-    "wheel",
-    "build",
+    "wheel>=0.38.1",
+    "build>=1",
 ]
 docs = [
-    "Sphinx>=7.0,<8.1",
+    "Sphinx>=7.0,<8.2",
     "sphinx-design>=0.4.0",
     "sphinx-autodoc-typehints>=1.22,<3.0",
     "sphinxcontrib-fulltoc~=1.2.0",
@@ -104,7 +104,7 @@ version-file = "netket/_version.py"
 [tool.uv]
 dev-dependencies = [
   "netket[dev]",
-  "mpi4py",
+  "mpi4py>=3.0.1",
   "jax>=0.4.33"
 ]