NetKet 3.14.2__tar.gz → 3.14.3__tar.gz

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

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

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

{netket-3.14.2 → netket-3.14.3}/netket/experimental/driver/tdvp_common.py

@@ -226,7 +226,9 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         if tstops is not None and (
             np.any(np.less(tstops, self.t)) or np.any(np.greater(tstops, t_end))
         ):
-            raise ValueError(f"All tstops must be in range [t, t + T]=[{self.t}, {T}]")
+            raise ValueError(
+                f"All tstops must be in range [t, t + T]=[{self.t}, {t_end}]"
+            )
 
         if tstops is not None and len(tstops) > 0:
             tstops = np.sort(tstops)

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

@@ -59,7 +59,6 @@ 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
 ):
@@ -71,6 +70,20 @@ 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,
+):
+    """
+    local_value kernel for MCState for jax-compatible operators
+    """
+
+    σ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)
+
+    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):
     """
@@ -183,10 +196,15 @@ def local_value_kernel_jax_chunked(
     """
     local_value kernel for MCState and jaxcoompatible operators
     """
-    local_value_kernel = lambda s: local_value_kernel_jax(logpsi, pars, s, O)
 
+    if chunk_size >= O.max_conn_size:
+        local_value_kernel = lambda s: local_value_kernel_jax(logpsi, pars, s, O)
+    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)
     )
 
     return local_value_chunked(σ)
+