NetKet 3.20.3__tar.gz → 3.20.4__tar.gz

{netket-3.20.3 → netket-3.20.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.20.3
+Version: 3.20.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.20.3 → netket-3.20.4}/netket/_src/vqs/transformed_vstate.py

@@ -42,6 +42,8 @@ def apply_operator(operator, vstate, *, seed=None):
         chunk_size = max(vstate.chunk_size // operator.max_conn_size, 1)
 
     if isinstance(vstate, FullSumState):
+        del seed
+
         transformed_vstate = FullSumState(
             hilbert=vstate.hilbert,
             apply_fun=transformed_apply_fun,
@@ -58,14 +60,18 @@ def apply_operator(operator, vstate, *, seed=None):
             n_samples=vstate.n_samples,
             n_discard_per_chain=vstate.n_discard_per_chain,
             chunk_size=chunk_size,
+            sampler_seed=seed,
         )
         if isinstance(vstate.sampler_state, MetropolisSamplerState):
+            # fold in a random counter so that we don't have the same seed as sampler
+            # itself.
+            seed = jax.random.fold_in(PRNGKey(seed), 123)
+
             x = vstate.sampler_state.σ
             xp, mels = operator.get_conn_padded(x)
-            seed = PRNGKey(seed)
             ids = jax.random.randint(seed, (x.shape[0],), 0, operator.max_conn_size)
             new_x = xp[jnp.arange(x.shape[0]), ids, :]
-            transformed_vstate.sampler_state = new_x
-
-        transformed_vstate.sampler_state = vstate.sampler_state.replace(σ=new_x)
+            transformed_vstate.sampler_state = transformed_vstate.sampler_state.replace(
+                σ=new_x
+            )
         return transformed_vstate

{netket-3.20.3 → netket-3.20.4}/netket/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.20.3'
-__version_tuple__ = version_tuple = (3, 20, 3)
+__version__ = version = '3.20.4'
+__version_tuple__ = version_tuple = (3, 20, 4)
 
 __commit_id__ = commit_id = None

{netket-3.20.3 → netket-3.20.4}/netket/errors.py

@@ -394,6 +394,114 @@ class JaxOperatorNotConvertibleToNumba(NetketError):
         )
 
 
+class JaxOperatorGetConnInJitError(NetketError):
+    """Illegal attempt to use :meth:`~netket.operator.DiscreteOperator.get_conn_flattened` or
+    similar methods inside of a Jax function transformation with a Jax operator.
+
+    This happens when calling :meth:`~netket.operator.DiscreteOperator.get_conn_flattened` or
+    :meth:`~netket.operator.DiscreteOperator.get_conn` inside of a function that is being
+    transformed by jax with transformations such as :func:`jax.jit`, :func:`jax.grad`,
+    :func:`jax.vmap`, or other jax function transformations.
+
+    These methods require materializing the output to a concrete numpy array, which is not
+    compatible with Jax's abstract tracing mechanism used in function transformations.
+
+    To avoid this error, use :meth:`~netket.operator.DiscreteOperator.get_conn_padded` instead,
+    which returns Jax arrays that are fully compatible with jax function transformations.
+
+    Example
+    -------
+
+    Instead of:
+
+    .. code-block:: python
+
+        import netket as nk
+        import jax
+        import jax.numpy as jnp
+
+        graph = nk.graph.Chain(10)
+        hilbert = nk.hilbert.Spin(1/2, graph.n_nodes)
+        ham = nk.operator.Ising(hilbert, graph, h=1.0)
+        ham_jax = ham.to_jax_operator()
+
+        ma = nk.models.RBM()
+        pars = ma.init(jax.random.PRNGKey(1), jnp.ones((2,graph.n_nodes)))
+
+        samples = hilbert.all_states()
+
+        @jax.jit
+        def compute_local_energies(pars, ham_jax, s):
+            # This will raise JaxOperatorGetConnInJitError
+            sp, mels = ham_jax.get_conn_flattened(s)
+
+            logpsi_sp = ma.apply(pars, sp)
+            logpsi_s = jnp.expand_dims(ma.apply(pars, s), -1)
+
+            return jnp.sum(mels * jnp.exp(logpsi_sp-logpsi_s), axis=-1)
+
+        elocs = compute_local_energies(pars, ham_jax, samples)
+
+    Use:
+
+    .. code-block:: python
+
+        import netket as nk
+        import jax
+        import jax.numpy as jnp
+
+        graph = nk.graph.Chain(10)
+        hilbert = nk.hilbert.Spin(1/2, graph.n_nodes)
+        ham = nk.operator.Ising(hilbert, graph, h=1.0)
+        ham_jax = ham.to_jax_operator()
+
+        ma = nk.models.RBM()
+        pars = ma.init(jax.random.PRNGKey(1), jnp.ones((2,graph.n_nodes)))
+
+        samples = hilbert.all_states()
+
+        @jax.jit
+        def compute_local_energies(pars, ham_jax, s):
+            # Use get_conn_padded instead - this works in jax.jit
+            sp, mels = ham_jax.get_conn_padded(s)
+
+            logpsi_sp = ma.apply(pars, sp)
+            logpsi_s = jnp.expand_dims(ma.apply(pars, s), -1)
+
+            return jnp.sum(mels * jnp.exp(logpsi_sp-logpsi_s), axis=-1)
+
+        elocs = compute_local_energies(pars, ham_jax, samples)
+
+    .. note::
+
+        :meth:`~netket.operator.DiscreteOperator.get_conn_padded` returns arrays with a fixed
+        padding dimension for the connected elements, which makes them compatible with jax
+        transformations. The padding is filled with zeros in the matrix elements array.
+
+    """
+
+    def __init__(self, operator):
+        operator_type = type(operator)
+        super().__init__(
+            f"\n"
+            f"Attempted to use 'get_conn_flattened' or 'get_conn' method of a Jax operator ({operator_type}) "
+            f"inside a Jax function transformation (jax.jit, jax.grad, jax.vmap & others)."
+            f"\n\n"
+            f"These methods require materializing the output to concrete numpy arrays, which is "
+            f"not compatible with Jax's abstract tracing mechanism."
+            f"\n\n"
+            f"To fix this error, use 'get_conn_padded' instead, which returns Jax arrays that "
+            f"are fully compatible with jax transformations:"
+            f"\n\n"
+            f"    # Instead of:"
+            f"\n    xp, mels = operator.get_conn(x)"
+            f"\n    xp, mels = operator.get_conn_flattened(x)"
+            f"\n\n"
+            f"    # Use:"
+            f"\n    xp, mels = operator.get_conn_padded(x)"
+        )
+
+
 #################################################
 # Jacobian and QGT errors                       #
 #################################################

{netket-3.20.3 → netket-3.20.4}/netket/operator/_discrete_operator_jax.py

@@ -20,6 +20,7 @@ import jax
 import jax.numpy as jnp
 from jax.experimental.sparse import JAXSparse, BCOO, BCSR
 
+from netket.errors import concrete_or_error, JaxOperatorGetConnInJitError
 from netket.operator import AbstractOperator, DiscreteOperator
 from netket.utils.optional_deps import import_optional_dependency
 
@@ -168,7 +169,13 @@ class DiscreteJaxOperator(DiscreteOperator):
 
         """
         xp, mels = self.get_conn_padded(x)
-        xp = np.array(xp)
+
+        xp = concrete_or_error(
+            np.asarray,
+            xp,
+            JaxOperatorGetConnInJitError,
+            self,
+        )
         mels = np.array(mels)
 
         if pad: