Repository: netket/netket

Issue:
LocalOperatorJax tracer leak (MetropolisHamiltonianJax breaks) 
cc @inailuig 

MWE:
```python
import netket as nk

L = 4
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
ha = nk.operator.IsingJax(hilbert=hi, graph=g, h=1.0).to_local_operator()
ma = nk.models.RBM(alpha=1, param_dtype=float)
sa = nk.sampler.MetropolisHamiltonian(hi, ha, n_chains=16)
vs = nk.vqs.MCState(sa, ma, n_samples=1008, n_discard_per_chain=10)
vs.sample()
```

error
```python
File ~/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:252, in LocalOperatorJax._get_conn_padded(self, x)
    249 shape = x.shape
    250 x = x.reshape(-1, x.shape[-1])
--> 252 xp, mels, n_conn = _local_operator_kernel_jax(
    253     self._nonzero_diagonal,
    254     self._max_conn_size,
    255     self._mel_cutoff,
    256     (
    257         self._local_states,
    258         self._acting_on,
    259         self._n_conns,
    260         self._diag_mels,
    261         self._x_prime,
    262         self._mels,
    263         self._basis,
    264         self._constant,
    265     ),
    266     x,
    267 )
    268 xp = xp.reshape(shape[:-1] + xp.shape[-2:])
    269 mels = mels.reshape(shape[:-1] + mels.shape[-1:])

    [... skipping hidden 5 frame]

File ~/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py:1720, in DynamicJaxprTracer._assert_live(self)
   1718 def _assert_live(self) -> None:
   1719   if not self._trace.main.jaxpr_stack:  # type: ignore
-> 1720     raise core.escaped_tracer_error(self, None)

UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[4,1,2] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
The function being traced when the value leaked was scanned_fun at /Users/filippo.vicentini/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/jax/_src/lax/control_flow/loops.py:1988 traced for scan.
------------------------------
The leaked intermediate value was created on line /Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/compile_helpers.py:232:41 (pack_internals_jax).
------------------------------
When the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
------------------------------
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/sampler/rules/hamiltonian.py:158:19 (HamiltonianRuleJax.transition)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:275:22 (LocalOperatorJax.get_conn_padded)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:247:8 (LocalOperatorJax._get_conn_padded)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:226:19 (LocalOperatorJax._setup)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/compile_helpers.py:232:41 (pack_internals_jax)
------------------------------

To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError
```

tracer leak checker suggests those tracers are created in `operator/_local_operator/compile_helpers.py:232:41 

```python
Traced<ShapedArray(int64[4,1])>with<DynamicJaxprTrace(level=3/0)>
This DynamicJaxprTracer was created on line /Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/compile_helpers.py:232:41 (pack_internals_jax)
<DynamicJaxprTracer 6432111440> is referred to by <list 6434208192>[0]
<list 6434208192> is referred to by <LocalOperatorJax 6417723152>._basis
<LocalOperatorJax 6417723152> is referred to by <dict 4401173760>['ha']
<dict 4401173760> is referred to by <TerminalInteractiveShell 4394007632>.user_ns
<TerminalInteractiveShell 4394007632> is referred to by <dict 4389887616>['_instance']
<dict 4389887616> is referred to by <MetaHasTraits 5016159744>
<MetaHasTraits 5016159744> is referred to by IPython.InteractiveShell
```
LocalOperatorJax tracer leak (MetropolisHamiltonianJax breaks) 
cc @inailuig 

MWE:
```python
import netket as nk

L = 4
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
ha = nk.operator.IsingJax(hilbert=hi, graph=g, h=1.0).to_local_operator()
ma = nk.models.RBM(alpha=1, param_dtype=float)
sa = nk.sampler.MetropolisHamiltonian(hi, ha, n_chains=16)
vs = nk.vqs.MCState(sa, ma, n_samples=1008, n_discard_per_chain=10)
vs.sample()
```

error
```python
File ~/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:252, in LocalOperatorJax._get_conn_padded(self, x)
    249 shape = x.shape
    250 x = x.reshape(-1, x.shape[-1])
--> 252 xp, mels, n_conn = _local_operator_kernel_jax(
    253     self._nonzero_diagonal,
    254     self._max_conn_size,
    255     self._mel_cutoff,
    256     (
    257         self._local_states,
    258         self._acting_on,
    259         self._n_conns,
    260         self._diag_mels,
    261         self._x_prime,
    262         self._mels,
    263         self._basis,
    264         self._constant,
    265     ),
    266     x,
    267 )
    268 xp = xp.reshape(shape[:-1] + xp.shape[-2:])
    269 mels = mels.reshape(shape[:-1] + mels.shape[-1:])

    [... skipping hidden 5 frame]

File ~/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py:1720, in DynamicJaxprTracer._assert_live(self)
   1718 def _assert_live(self) -> None:
   1719   if not self._trace.main.jaxpr_stack:  # type: ignore
-> 1720     raise core.escaped_tracer_error(self, None)

UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[4,1,2] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
The function being traced when the value leaked was scanned_fun at /Users/filippo.vicentini/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/jax/_src/lax/control_flow/loops.py:1988 traced for scan.
------------------------------
The leaked intermediate value was created on line /Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/compile_helpers.py:232:41 (pack_internals_jax).
------------------------------
When the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
------------------------------
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/sampler/rules/hamiltonian.py:158:19 (HamiltonianRuleJax.transition)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:275:22 (LocalOperatorJax.get_conn_padded)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:247:8 (LocalOperatorJax._get_conn_padded)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/jax.py:226:19 (LocalOperatorJax._setup)
/Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/compile_helpers.py:232:41 (pack_internals_jax)
------------------------------

To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError
```

tracer leak checker suggests those tracers are created in `operator/_local_operator/compile_helpers.py:232:41 

```python
Traced<ShapedArray(int64[4,1])>with<DynamicJaxprTrace(level=3/0)>
This DynamicJaxprTracer was created on line /Users/filippo.vicentini/Dropbox/Ricerca/Codes/Python/netket/netket/operator/_local_operator/compile_helpers.py:232:41 (pack_internals_jax)
<DynamicJaxprTracer 6432111440> is referred to by <list 6434208192>[0]
<list 6434208192> is referred to by <LocalOperatorJax 6417723152>._basis
<LocalOperatorJax 6417723152> is referred to by <dict 4401173760>['ha']
<dict 4401173760> is referred to by <TerminalInteractiveShell 4394007632>.user_ns
<TerminalInteractiveShell 4394007632> is referred to by <dict 4389887616>['_instance']
<dict 4389887616> is referred to by <MetaHasTraits 5016159744>
<MetaHasTraits 5016159744> is referred to by IPython.InteractiveShell
```


Discussion:
I think the leak is because we call `_setup()` in the jaxoperator inside the metropolis rule (`sampler.rule.operator`) inside of the sampling loop of metropolis, but then this object 'leaks out' because we keep the operator around. 

I think the leak is because we call `_setup()` in the jaxoperator inside the metropolis rule (`sampler.rule.operator`) inside of the sampling loop of metropolis, but then this object 'leaks out' because we keep the operator around. 


Resolving patch:
diff --git a/netket/sampler/metropolis.py b/netket/sampler/metropolis.py
--- a/netket/sampler/metropolis.py
+++ b/netket/sampler/metropolis.py
@@ -695,7 +695,7 @@ def MetropolisHamiltonian(hilbert, hamiltonian, **kwargs) -> MetropolisSampler:
        >>> # Construct a MetropolisHamiltonian Sampler
        >>> sa = nk.sampler.MetropolisHamiltonian(hi, hamiltonian=ha)
        >>> print(sa)
-       MetropolisSampler(rule = HamiltonianRuleNumba(Ising(J=1.0, h=1.0; dim=100)), n_chains = 16, sweep_size = 100, reset_chains = False, machine_power = 2, dtype = <class 'float'>)
+       MetropolisSampler(rule = HamiltonianRuleNumba(operator=Ising(J=1.0, h=1.0; dim=100)), n_chains = 16, sweep_size = 100, reset_chains = False, machine_power = 2, dtype = <class 'float'>)
     """
     from .rules import HamiltonianRule
 
diff --git a/netket/sampler/rules/hamiltonian.py b/netket/sampler/rules/hamiltonian.py
--- a/netket/sampler/rules/hamiltonian.py
+++ b/netket/sampler/rules/hamiltonian.py
@@ -23,7 +23,7 @@
 from numba import jit
 from numba4jax import njit4jax
 
-from netket.operator import AbstractOperator, DiscreteJaxOperator
+from netket.operator import DiscreteOperator, DiscreteJaxOperator
 from netket.utils import struct
 from netket.jax.sharding import sharding_decorator, with_samples_sharding_constraint
 
@@ -35,50 +35,51 @@ class HamiltonianRuleBase(MetropolisRule):
     Rule proposing moves according to the terms in an operator.
     """
 
-    operator: AbstractOperator = struct.field(pytree_node=False)
+    # operator: AbstractOperator = struct.field(pytree_node=False / True depending on jax or not jax)
     """The (hermitian) operator giving the transition amplitudes."""
 
-    def __init__(self, operator: AbstractOperator):
-        self.operator = operator
-
-    def init_state(rule, sampler, machine, params, key):
-        if sampler.hilbert != rule.operator.hilbert:
+    def init_state(self, sampler, machine, params, key):
+        if sampler.hilbert != self.operator.hilbert:
             raise ValueError(
                 f"""
             The hilbert space of the sampler ({sampler.hilbert}) and the hilbert space
-            of the operator ({rule.operator.hilbert}) for HamiltonianRule must be the same.
+            of the operator ({self.operator.hilbert}) for HamiltonianRule must be the same.
             """
             )
         return super().init_state(sampler, machine, params, key)
 
-    def __post_init__(self):
-        # Raise errors if hilbert is not an Hilbert
-        if not isinstance(self.operator, AbstractOperator):
-            raise TypeError(
-                "Argument to HamiltonianRule must be a valid operator, "
-                f"but operator is a {type(self.operator)}."
-            )
-
-    def __repr__(self):
-        return f"HamiltonianRuleNumba({self.operator})"
-
 
 @struct.dataclass
 class HamiltonianRuleNumba(HamiltonianRuleBase):
-    """
+    r"""
     Rule proposing moves according to the terms in an operator.
 
     In this case, the transition matrix is taken to be:
 
     .. math::
 
-       T( \\mathbf{s} \\rightarrow \\mathbf{s}^\\prime) = \\frac{1}{\\mathcal{N}(\\mathbf{s})}\\theta(|H_{\\mathbf{s},\\mathbf{s}^\\prime}|),
+       T( \mathbf{s} \rightarrow \mathbf{s}^\\prime) = \frac{1}{\mathcal{N}(\mathbf{s})}\theta(|H_{\mathbf{s},\mathbf{s}^\prime}|),
 
     This rule only works on CPU! If you want to use it on GPU, you
     must use the numpy variant :class:`netket.sampler.rules.HamiltonianRuleNumpy`
     together with the numpy metropolis sampler :class:`netket.sampler.MetropolisSamplerNumpy`.
     """
 
+    operator: DiscreteOperator = struct.field(pytree_node=False)
+    """The (hermitian) operator giving the transition amplitudes."""
+
+    def __init__(self, operator: DiscreteOperator):
+        if not isinstance(operator, DiscreteOperator):
+            raise TypeError(
+                "Argument to HamiltonianRule must be a valid operator, "
+                f"but operator is a {type(operator)}."
+            )
+        # call _setup on the operator if it exists, to warmup the cache and
+        # avoid calling it in a numba callback which might break things.
+        if hasattr(operator, "_setup"):
+            operator._setup()
+        self.operator = operator
+
     def transition(rule, sampler, machine, parameters, state, key, σ):
         """
         This implements the transition rule for `DiscreteOperator`s that are
@@ -126,9 +127,6 @@ def __transition(args):
 
         return σp, log_prob_correction
 
-    def __repr__(self):
-        return f"HamiltonianRuleNumba({self.operator})"
-
 
 @jit(nopython=True)
 def _choose(vp, sections, rand_vec, out, w):
@@ -142,18 +140,29 @@ def _choose(vp, sections, rand_vec, out, w):
 
 @struct.dataclass
 class HamiltonianRuleJax(HamiltonianRuleBase):
-    """
+    r"""
     Rule proposing moves according to the terms in an operator.
 
     In this case, the transition matrix is taken to be:
 
     .. math::
 
-       T( \\mathbf{s} \\rightarrow \\mathbf{s}^\\prime) = \\frac{1}{\\mathcal{N}(\\mathbf{s})}\\theta(|H_{\\mathbf{s},\\mathbf{s}^\\prime}|),
+       T( \mathbf{s} \rightarrow \mathbf{s}^\prime) = \frac{1}{\mathcal{N}(\mathbf{s})}\theta(|H_{\mathbf{s},\mathbf{s}^\prime}|),
 
     This rule only works with operators which are written in jax.
     """
 
+    operator: DiscreteJaxOperator = struct.field(pytree_node=True)
+    """The (hermitian) operator giving the transition amplitudes."""
+
+    def __init__(self, operator: DiscreteJaxOperator):
+        if not isinstance(operator, DiscreteJaxOperator):
+            raise TypeError(
+                "Argument to HamiltonianRule must be a valid operator, "
+                f"but operator is a {type(operator)}."
+            )
+        self.operator = operator
+
     def transition(self, _0, _1, _2, _3, key, x):
         xp, mels = self.operator.get_conn_padded(x)
 
@@ -202,9 +211,6 @@ def transition(self, _0, _1, _2, _3, key, x):
 
         return x_proposed, log_prob_corr
 
-    def __repr__(self):
-        return f"HamiltonianRuleJax({self.operator})"
-
 
 def HamiltonianRule(operator):
     r"""


Test patch:
diff --git a/test/sampler/test_sampler.py b/test/sampler/test_sampler.py
--- a/test/sampler/test_sampler.py
+++ b/test/sampler/test_sampler.py
@@ -603,3 +603,33 @@ def test_multiplerules_pt(model_and_weights):
         chain_length=10,
     )
     assert samples.shape == (sa.n_chains, 10, hi.size)
+
+
+def test_hamiltonian_jax_sampler_isleaf():
+    g = nk.graph.Hypercube(length=4, n_dim=1, pbc=True)
+
+    hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
+    rule1 = nk.sampler.rules.HamiltonianRule(
+        nk.operator.IsingJax(hilbert=hi, graph=g, h=1.0)
+    )
+    rule2 = nk.sampler.rules.HamiltonianRule(
+        nk.operator.IsingJax(hilbert=hi, graph=g, h=1.0)
+    )
+    leaf1, struct1 = jax.tree_util.tree_flatten(rule1)
+    leaf2, struct2 = jax.tree_util.tree_flatten(rule2)
+
+    # if the structures are identical, the operators must have been unpacked.
+    assert struct1 == struct2
+    assert hash(struct1) == hash(struct2)
+
+    # check contained in leafs (this only works because the arrays are identically the same, but it
+    # is enough for this check):
+    for leaf in jax.tree_util.tree_leaves(rule1.operator):
+        found = False
+        for l in leaf1:
+            if leaf is l:
+                found = True
+                break
+        # If this fails, it is either because the operator is not a leaf ot the rule, or because jax changed
+        # some internals and the flattening does not return identical arrays anymore.
+        assert found