Repository: netket/netket

Issue:
Incomprehensible bug from PauliOperatorJax
It's a spooky action at a distance when converting to Pauli strings and then to jax operators...
The error is unreadable because it happens in the depth of jax jit calling logic, only after calling the first one expect...

Identified by @adrien-kahn

```python
import numpy as np
import jax
import jax.numpy as jnp
import netket as nk

# Define the Hilbert space
graph = nk.graph.Chain(4, pbc=True)
hi = nk.hilbert.Qubit(graph.n_nodes)

# Define the model
ma = nk.models.RBM()

# Define the MC variational state
sampler = nk.sampler.ExactSampler(hi)
vs = nk.vqs.MCState(sampler, ma)

# Define the Hamiltonian

hamiltonian = nk.operator.Heisenberg(hi, graph)
hamiltonian_pauli2 = hamiltonian.to_pauli_strings()
hamiltonian_pauli1 = hamiltonian.to_pauli_strings()
hamiltonian_jax2 = hamiltonian_pauli2.to_jax_operator()
hamiltonian_jax1 = hamiltonian_pauli1.to_jax_operator()

vs.expect(hamiltonian_jax1)
vs.expect(hamiltonian_jax2)
```

error
```python
Cell In[6], line 26
     23 hamiltonian_jax1 = hamiltonian_pauli1.to_jax_operator()
     25 mc_state.expect(hamiltonian_jax1)
---> 26 mc_state.expect(hamiltonian_jax2)

File ~/Dropbox/Ricerca/Codes/Python/netket/netket/vqs/mc/mc_state/state.py:585, in MCState.expect(self, O)
    569 def expect(self, O: AbstractOperator) -> Stats:
    570     r"""Estimates the quantum expectation value for a given operator
    571     :math:`O` or generic observable.
    572     In the case of a pure state :math:`\psi` and an operator, this is
   (...)
    583         :math:`\langle O\rangle`.
    584     """
--> 585     return expect(self, O, self.chunk_size)

    [... skipping hidden 1 frame]

File ~/Dropbox/Ricerca/Codes/Python/netket/netket/vqs/mc/mc_state/expect.py:111, in expect(vstate, Ô, chunk_size)
    108 σ, args = get_local_kernel_arguments(vstate, Ô)
    109 local_estimator_fun = get_local_kernel(vstate, Ô)
--> 111 res =  _expect(
    112     args,
    113 )
    114 print("got result")
    115 return res

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
```
Incomprehensible bug from PauliOperatorJax
It's a spooky action at a distance when converting to Pauli strings and then to jax operators...
The error is unreadable because it happens in the depth of jax jit calling logic, only after calling the first one expect...

Identified by @adrien-kahn

```python
import numpy as np
import jax
import jax.numpy as jnp
import netket as nk

# Define the Hilbert space
graph = nk.graph.Chain(4, pbc=True)
hi = nk.hilbert.Qubit(graph.n_nodes)

# Define the model
ma = nk.models.RBM()

# Define the MC variational state
sampler = nk.sampler.ExactSampler(hi)
vs = nk.vqs.MCState(sampler, ma)

# Define the Hamiltonian

hamiltonian = nk.operator.Heisenberg(hi, graph)
hamiltonian_pauli2 = hamiltonian.to_pauli_strings()
hamiltonian_pauli1 = hamiltonian.to_pauli_strings()
hamiltonian_jax2 = hamiltonian_pauli2.to_jax_operator()
hamiltonian_jax1 = hamiltonian_pauli1.to_jax_operator()

vs.expect(hamiltonian_jax1)
vs.expect(hamiltonian_jax2)
```

error
```python
Cell In[6], line 26
     23 hamiltonian_jax1 = hamiltonian_pauli1.to_jax_operator()
     25 mc_state.expect(hamiltonian_jax1)
---> 26 mc_state.expect(hamiltonian_jax2)

File ~/Dropbox/Ricerca/Codes/Python/netket/netket/vqs/mc/mc_state/state.py:585, in MCState.expect(self, O)
    569 def expect(self, O: AbstractOperator) -> Stats:
    570     r"""Estimates the quantum expectation value for a given operator
    571     :math:`O` or generic observable.
    572     In the case of a pure state :math:`\psi` and an operator, this is
   (...)
    583         :math:`\langle O\rangle`.
    584     """
--> 585     return expect(self, O, self.chunk_size)

    [... skipping hidden 1 frame]

File ~/Dropbox/Ricerca/Codes/Python/netket/netket/vqs/mc/mc_state/expect.py:111, in expect(vstate, Ô, chunk_size)
    108 σ, args = get_local_kernel_arguments(vstate, Ô)
    109 local_estimator_fun = get_local_kernel(vstate, Ô)
--> 111 res =  _expect(
    112     args,
    113 )
    114 print("got result")
    115 return res

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
```


Discussion:

Resolving patch:
diff --git a/netket/operator/_pauli_strings/jax.py b/netket/operator/_pauli_strings/jax.py
--- a/netket/operator/_pauli_strings/jax.py
+++ b/netket/operator/_pauli_strings/jax.py
@@ -22,8 +22,9 @@
 from jax.tree_util import register_pytree_node_class
 
 from netket.hilbert import AbstractHilbert, HomogeneousHilbert
-from netket.utils.types import DType
 from netket.errors import concrete_or_error, JaxOperatorSetupDuringTracingError
+from netket.utils.types import DType
+from netket.utils import HashableArray
 
 from .._discrete_operator_jax import DiscreteJaxOperator
 
@@ -382,11 +383,10 @@ def get_conn_padded(self, x):
 
     def tree_flatten(self):
         self._setup()
-
         data = (self.weights, self._x_flip_masks_stacked, self._z_data)
         metadata = {
             "hilbert": self.hilbert,
-            "operators": self.operators,
+            "operators": HashableArray(self.operators),
             "dtype": self.dtype,
         }
         return data, metadata
@@ -395,10 +395,12 @@ def tree_flatten(self):
     def tree_unflatten(cls, metadata, data):
         (weights, xm, zd) = data
         hi = metadata["hilbert"]
-        operators = metadata["operators"]
+        operators = metadata["operators"].wrapped
         dtype = metadata["dtype"]
 
-        op = cls(hi, operators, weights, dtype=dtype)
+        op = cls(hi, dtype=dtype)
+        op._operators = operators
+        op._weights = weights
         op._x_flip_masks_stacked = xm
         op._z_data = zd
         op._initialized = True


Test patch:
diff --git a/test/operator/test_operator.py b/test/operator/test_operator.py
--- a/test/operator/test_operator.py
+++ b/test/operator/test_operator.py
@@ -404,3 +404,28 @@ def _get_conn_padded(s):
 
     with pytest.raises(NumbaOperatorGetConnDuringTracingError):
         _get_conn_padded(state)
+
+
+def test_pauli_string_operators_hashable_pytree():
+    # Define the Hilbert space
+    graph = nk.graph.Chain(4, pbc=True)
+    hi = nk.hilbert.Qubit(graph.n_nodes)
+
+    # Define the model
+    ma = nk.models.RBM()
+
+    # Define the MC variational state
+    sampler = nk.sampler.ExactSampler(hi)
+    vs = nk.vqs.MCState(sampler, ma)
+
+    # Define the Hamiltonian
+
+    ha = nk.operator.Heisenberg(hi, graph)
+    hap2 = ha.to_pauli_strings()
+    hap1 = ha.to_pauli_strings()
+    haj2 = hap2.to_jax_operator()
+    haj1 = hap1.to_jax_operator()
+
+    e1 = vs.expect(haj1)
+    e2 = vs.expect(haj2)
+    jax.tree_map(np.testing.assert_allclose, e1, e2)