Repository: netket/netket

Issue:
Error with VMC_SRt
### Discussed in https://github.com/orgs/netket/discussions/1642

<div type='discussions-op-text'>

<sup>Originally posted by **aashmore** November  8, 2023</sup>
Hello!

First of all, let me say thank you to all of the NetKet developers. What an amazing tool!

I'm trying to get VMC_SRt working (amazing that it has been added!). Unfortunately, it is currently giving me an error when I run the driver, while using standard SR works just fine.

The model is built on a discrete spin Hilbert space with complex parameters and complex output.

Using 
```python
gs = nk.driver.VMC(H, optimizer, variational_state=vstate, preconditioner=nk.optimizer.SR(diag_shift=0.1, holomorphic=False), holomorphic=False)
```
the model trains and converges with no errors (and reproduces known results from the literature).

Changing this to 
```python
import netket.experimental as nkx
gs = nkx.driver.VMC_SRt(H, optimizer, diag_shift=0.01, variational_state=vstate, jacobian_mode="complex")
```
leads to the error
```python
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-33-d2c8dffbbffd>](https://localhost:8080/#) in <cell line: 3>()
      1 log = nk.logging.RuntimeLog()
      2 
----> 3 gs.run(n_iter=100, out=log)
      4 
      5 ffn_energy = vstate.expect(H)

8 frames
[/usr/local/lib/python3.10/dist-packages/netket/driver/abstract_variational_driver.py](https://localhost:8080/#) in run(self, n_iter, out, obs, show_progress, save_params_every, write_every, step_size, callback)
    257             first_step = True
    258 
--> 259             for step in self.iter(n_iter, step_size):
    260                 log_data = self.estimate(obs)
    261                 self._log_additional_data(log_data, step)

[/usr/local/lib/python3.10/dist-packages/netket/driver/abstract_variational_driver.py](https://localhost:8080/#) in iter(self, n_steps, step)
    167         for _ in range(0, n_steps, step):
    168             for i in range(0, step):
--> 169                 dp = self._forward_and_backward()
    170                 if i == 0:
    171                     yield self.step_count

[/usr/local/lib/python3.10/dist-packages/netket/experimental/driver/vmc_srt.py](https://localhost:8080/#) in _forward_and_backward(self)
    244         )
    245 
--> 246         self._dp = self._unravel_params_fn(updates)
    247 
    248         return self._dp

    [... skipping hidden 13 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/flatten_util.py](https://localhost:8080/#) in unravel_pytree(treedef, unravel_list, flat)
     51 
     52 def unravel_pytree(treedef, unravel_list, flat):
---> 53   return tree_unflatten(treedef, unravel_list(flat))
     54 
     55 def _ravel_list(lst):

    [... skipping hidden 1 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/flatten_util.py](https://localhost:8080/#) in _unravel_list_single_dtype(indices, shapes, arr)
     76 def _unravel_list_single_dtype(indices, shapes, arr):
     77   chunks = jnp.split(arr, indices[:-1])
---> 78   return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
     79 
     80 def _unravel_list(indices, shapes, from_dtypes, to_dtype, arr):

[/usr/local/lib/python3.10/dist-packages/jax/_src/flatten_util.py](https://localhost:8080/#) in <listcomp>(.0)
     76 def _unravel_list_single_dtype(indices, shapes, arr):
     77   chunks = jnp.split(arr, indices[:-1])
---> 78   return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
     79 
     80 def _unravel_list(indices, shapes, from_dtypes, to_dtype, arr):

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py](https://localhost:8080/#) in meth(self, *args, **kwargs)
    731 def _forward_method_to_aval(name):
    732   def meth(self, *args, **kwargs):
--> 733     return getattr(self.aval, name).fun(self, *args, **kwargs)
    734   return meth
    735 

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py](https://localhost:8080/#) in _reshape(a, order, *args)
    141   newshape = _compute_newshape(a, args[0] if len(args) == 1 else args)
    142   if order == "C":
--> 143     return lax.reshape(a, newshape, None)
    144   elif order == "F":
    145     dims = list(range(a.ndim)[::-1])

    [... skipping hidden 8 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/lax/lax.py](https://localhost:8080/#) in _reshape_shape_rule(operand, new_sizes, dimensions)
   3376       not math.prod(np.shape(operand)) == math.prod(new_sizes)):
   3377     msg = 'reshape total size must be unchanged, got new_sizes {} for shape {}.'
-> 3378     raise TypeError(msg.format(new_sizes, np.shape(operand)))
   3379   if dimensions is not None:
   3380     if set(dimensions) != set(range(np.ndim(operand))):

TypeError: reshape total size must be unchanged, got new_sizes (2, 3, 7) for shape (84,).
```

From my (amateur) reading of the trace, it looks like something is going wrong with chunking? The model itself is rather complicated, so if it isn't a problem with VMC_SRt itself, I'll put together a minimal working example and edit my question. I'm just a little confused as it why it would work with standard VMC but not VMC_SRt, since my understand is that the difference is just in how one computes a certain matrix.

Thank you in advance!</div>
Error with VMC_SRt
### Discussed in https://github.com/orgs/netket/discussions/1642

<div type='discussions-op-text'>

<sup>Originally posted by **aashmore** November  8, 2023</sup>
Hello!

First of all, let me say thank you to all of the NetKet developers. What an amazing tool!

I'm trying to get VMC_SRt working (amazing that it has been added!). Unfortunately, it is currently giving me an error when I run the driver, while using standard SR works just fine.

The model is built on a discrete spin Hilbert space with complex parameters and complex output.

Using 
```python
gs = nk.driver.VMC(H, optimizer, variational_state=vstate, preconditioner=nk.optimizer.SR(diag_shift=0.1, holomorphic=False), holomorphic=False)
```
the model trains and converges with no errors (and reproduces known results from the literature).

Changing this to 
```python
import netket.experimental as nkx
gs = nkx.driver.VMC_SRt(H, optimizer, diag_shift=0.01, variational_state=vstate, jacobian_mode="complex")
```
leads to the error
```python
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-33-d2c8dffbbffd>](https://localhost:8080/#) in <cell line: 3>()
      1 log = nk.logging.RuntimeLog()
      2 
----> 3 gs.run(n_iter=100, out=log)
      4 
      5 ffn_energy = vstate.expect(H)

8 frames
[/usr/local/lib/python3.10/dist-packages/netket/driver/abstract_variational_driver.py](https://localhost:8080/#) in run(self, n_iter, out, obs, show_progress, save_params_every, write_every, step_size, callback)
    257             first_step = True
    258 
--> 259             for step in self.iter(n_iter, step_size):
    260                 log_data = self.estimate(obs)
    261                 self._log_additional_data(log_data, step)

[/usr/local/lib/python3.10/dist-packages/netket/driver/abstract_variational_driver.py](https://localhost:8080/#) in iter(self, n_steps, step)
    167         for _ in range(0, n_steps, step):
    168             for i in range(0, step):
--> 169                 dp = self._forward_and_backward()
    170                 if i == 0:
    171                     yield self.step_count

[/usr/local/lib/python3.10/dist-packages/netket/experimental/driver/vmc_srt.py](https://localhost:8080/#) in _forward_and_backward(self)
    244         )
    245 
--> 246         self._dp = self._unravel_params_fn(updates)
    247 
    248         return self._dp

    [... skipping hidden 13 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/flatten_util.py](https://localhost:8080/#) in unravel_pytree(treedef, unravel_list, flat)
     51 
     52 def unravel_pytree(treedef, unravel_list, flat):
---> 53   return tree_unflatten(treedef, unravel_list(flat))
     54 
     55 def _ravel_list(lst):

    [... skipping hidden 1 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/flatten_util.py](https://localhost:8080/#) in _unravel_list_single_dtype(indices, shapes, arr)
     76 def _unravel_list_single_dtype(indices, shapes, arr):
     77   chunks = jnp.split(arr, indices[:-1])
---> 78   return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
     79 
     80 def _unravel_list(indices, shapes, from_dtypes, to_dtype, arr):

[/usr/local/lib/python3.10/dist-packages/jax/_src/flatten_util.py](https://localhost:8080/#) in <listcomp>(.0)
     76 def _unravel_list_single_dtype(indices, shapes, arr):
     77   chunks = jnp.split(arr, indices[:-1])
---> 78   return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
     79 
     80 def _unravel_list(indices, shapes, from_dtypes, to_dtype, arr):

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py](https://localhost:8080/#) in meth(self, *args, **kwargs)
    731 def _forward_method_to_aval(name):
    732   def meth(self, *args, **kwargs):
--> 733     return getattr(self.aval, name).fun(self, *args, **kwargs)
    734   return meth
    735 

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py](https://localhost:8080/#) in _reshape(a, order, *args)
    141   newshape = _compute_newshape(a, args[0] if len(args) == 1 else args)
    142   if order == "C":
--> 143     return lax.reshape(a, newshape, None)
    144   elif order == "F":
    145     dims = list(range(a.ndim)[::-1])

    [... skipping hidden 8 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/lax/lax.py](https://localhost:8080/#) in _reshape_shape_rule(operand, new_sizes, dimensions)
   3376       not math.prod(np.shape(operand)) == math.prod(new_sizes)):
   3377     msg = 'reshape total size must be unchanged, got new_sizes {} for shape {}.'
-> 3378     raise TypeError(msg.format(new_sizes, np.shape(operand)))
   3379   if dimensions is not None:
   3380     if set(dimensions) != set(range(np.ndim(operand))):

TypeError: reshape total size must be unchanged, got new_sizes (2, 3, 7) for shape (84,).
```

From my (amateur) reading of the trace, it looks like something is going wrong with chunking? The model itself is rather complicated, so if it isn't a problem with VMC_SRt itself, I'll put together a minimal working example and edit my question. I'm just a little confused as it why it would work with standard VMC but not VMC_SRt, since my understand is that the difference is just in how one computes a certain matrix.

Thank you in advance!</div>


Discussion:

Resolving patch:
diff --git a/netket/experimental/driver/vmc_srt.py b/netket/experimental/driver/vmc_srt.py
--- a/netket/experimental/driver/vmc_srt.py
+++ b/netket/experimental/driver/vmc_srt.py
@@ -20,7 +20,9 @@
 
 
 @partial(jax.jit, static_argnames=("mode", "solver_fn"))
-def SRt(O_L, local_energies, diag_shift, *, mode, solver_fn, e_mean=None):
+def SRt(
+    O_L, local_energies, diag_shift, *, mode, solver_fn, e_mean=None, params_structure
+):
     """
     For more details, see `https://arxiv.org/abs/2310.05715'. In particular,
     the following parallel implementation is described in Appendix "Distributed SR computation".
@@ -41,7 +43,10 @@ def SRt(O_L, local_energies, diag_shift, *, mode, solver_fn, e_mean=None):
     dv = -2.0 * de / N_mc**0.5
 
     if mode == "complex":
-        O_L = jnp.concatenate((O_L[:, 0], O_L[:, 1]), axis=0)
+        # Concatenate the real and imaginary derivatives of the ansatz
+        # O_L = jnp.concatenate((O_L[:, 0], O_L[:, 1]), axis=0)
+        O_L = jnp.transpose(O_L, (1, 0, 2)).reshape(-1, O_L.shape[-1])
+
         dv = jnp.concatenate((jnp.real(dv), -jnp.imag(dv)), axis=-1)
     elif mode == "real":
         dv = dv.real
@@ -77,6 +82,12 @@ def SRt(O_L, local_energies, diag_shift, *, mode, solver_fn, e_mean=None):
     updates = O_L.T @ aus_vector
     updates, token = mpi.mpi_allreduce_sum_jax(updates, token=token)
 
+    # If complex mode and we have complex parameters, we need
+    # To repack the real coefficients in order to get complex updates
+    if mode == "complex" and nkjax.tree_leaf_iscomplex(params_structure):
+        np = updates.shape[-1] // 2
+        updates = (updates[:np] + 1j * updates[np:]) / 2
+
     return -updates
 
 
@@ -143,14 +154,6 @@ def __init__(
         """
         super().__init__(hamiltonian, optimizer, variational_state=variational_state)
 
-        self._ham = hamiltonian.collect()  # type: AbstractOperator
-        self.diag_shift = diag_shift
-        self.jacobian_mode = jacobian_mode
-        self._linear_solver_fn = linear_solver_fn
-
-        _, unravel_params_fn = ravel_pytree(self.state.parameters)
-        self._unravel_params_fn = jax.jit(unravel_params_fn)
-
         if self.state.n_parameters % mpi.n_nodes != 0:
             raise NotImplementedError(
                 f"""
@@ -173,6 +176,24 @@ def __init__(
                 stacklevel=2,
             )
 
+        self._ham = hamiltonian.collect()  # type: AbstractOperator
+        self.diag_shift = diag_shift
+        self.jacobian_mode = jacobian_mode
+        self._linear_solver_fn = linear_solver_fn
+
+        self._params_structure = jax.tree_map(
+            lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), self.state.parameters
+        )
+        if not nkjax.tree_ishomogeneous(self._params_structure):
+            raise ValueError(
+                "SRt only supports neural networks with all real or all complex "
+                "parameters. Hybrid structures are not yet supported (but we would welcome "
+                "contributions. Get in touch with us!)"
+            )
+
+        _, unravel_params_fn = ravel_pytree(self.state.parameters)
+        self._unravel_params_fn = jax.jit(unravel_params_fn)
+
     @property
     def jacobian_mode(self) -> str:
         """
@@ -228,7 +249,7 @@ def _forward_and_backward(self):
             mode=self.jacobian_mode,
             dense=True,
             center=True,
-        )  # * jaxcobians is centered
+        )  # jacobians is centered
 
         diag_shift = self.diag_shift
         if callable(self.diag_shift):
@@ -241,6 +262,7 @@ def _forward_and_backward(self):
             mode=self.jacobian_mode,
             solver_fn=self._linear_solver_fn,
             e_mean=self._loss_stats.Mean,
+            params_structure=self._params_structure,
         )
 
         self._dp = self._unravel_params_fn(updates)
diff --git a/netket/jax/_jacobian/default_mode.py b/netket/jax/_jacobian/default_mode.py
--- a/netket/jax/_jacobian/default_mode.py
+++ b/netket/jax/_jacobian/default_mode.py
@@ -56,7 +56,7 @@ def __eq__(self, o):
 HolomorphicMode = JacobianMode("holomorphic")
 
 
-@partial(jax.jit, static_argnames=("apply_fun", "holomorphic"))
+@partial(jax.jit, static_argnames=("apply_fun", "holomorphic", "warn"))
 def jacobian_default_mode(
     apply_fun: Callable[[PyTree, Array], Array],
     pars: PyTree,
@@ -67,7 +67,7 @@ def jacobian_default_mode(
     warn: bool = True,
 ) -> JacobianMode:
     """
-    Returns the default `mode` for {func}`nk.jax.jacobian` given a certain
+    Returns the default `mode` for {func}`netket.jax.jacobian` given a certain
     wave-function ansatz.
 
     This function uses an abstract evaluation of the ansatz to determine if
diff --git a/netket/jax/_jacobian/logic.py b/netket/jax/_jacobian/logic.py
--- a/netket/jax/_jacobian/logic.py
+++ b/netket/jax/_jacobian/logic.py
@@ -103,7 +103,7 @@ def jacobian(
             read the detailed discussion below.
         pdf: Optional coefficient that is used to multiply every row of the Jacobian.
             When performing calculations in full-summation, this can be used to
-            multiply every row by :math:`\abs{\psi(\sigma)}^2`, which is needed to
+            multiply every row by :math:`|\psi(\sigma)|^2`, which is needed to
             compute the correct average.
         chunk_size: Optional integer specifying the maximum number of samples for
             which the gradient is simulataneously computed. Low-values will
@@ -187,11 +187,12 @@ def jacobian(
 
     .. math::
 
-       O_k(\sigma) = \frac{\partial \ln\Re[\Psi(\sigma)]}{\partial \Re[\theta_k]}
+       O^{r}_k(\sigma) = \frac{\partial \ln\Re[\Psi(\sigma)]}{\partial \theta_k}
        \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,
-       O_k(\sigma) = \frac{\partial \ln\Im[\Psi(\sigma)]}{\partial \Re[\theta_k]}
+       O^{i}_k(\sigma) = \frac{\partial \ln\Im[\Psi(\sigma)]}{\partial \theta_k}
 
-    properly concatenated in a single PyTree for every set of parameters. In practice,
+    where :math:`O^{r}_k(\sigma)` and :math:`O^{i}_k(\sigma)` are real-valued pytrees
+    with the same shape as the original parameters. In practice,
     it should return a result roughly equivalent to the following listing:
 
     .. code:: python
@@ -207,17 +208,41 @@ def jacobian(
     do this for performance reason, but the downstream user is free to do it if
     he wishes.
 
-    **If some parameters** :math:`\theta_k` **are complex**, this mode returns the
-    derivatives of the real and imaginary part of the function,
+    If you wish to get the complex jacobian in the case of real parameters, it is
+    possible to define
 
     .. math::
 
-       O_k(\sigma) = \frac{\partial \ln\Re[\Psi(\sigma)]}{\partial \Re[\theta_k]}
+         O_k(\sigma) =  O^{r}_k(\sigma) + i O^{i}_k(\sigma)
+
+    which is now complex-valued. In code, this is equivalent to
+
+    .. code:: python
+
+      O_k_cmplx = jax.tree_map(lambda jri: jri[:, 0, :] + 1j* jri[:, 1, :], O_k)
+
+
+    **If some parameters** :math:`\theta_k` **are complex**, this mode splits the
+    :math:`N` complex parameters into :math:`2N` real parameters, where the first
+    block of :math:`N` parameters correspond to the real parts and the latter block
+    to the imaginary part, and then follows the logic discussed above.
+
+    In formulas, this can be seen as defining the vector of :math:`2N` real parameters
+
+    .. math::
+
+        \tilde\theta = (\Re[\theta], \Im[\theta])
+
+    and then computing the same quantities as above
+
+    .. math::
+
+       O^{r}_k(\sigma) = \frac{\partial \ln\Re[\Psi(\sigma)]}{\partial \tilde\theta_k]}
        \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,
-       O_k(\sigma) = \frac{\partial \ln\Im[\Psi(\sigma)]}{\partial \Re[\theta_k]}
+       O^{i}_k(\sigma) = \frac{\partial \ln\Im[\Psi(\sigma)]}{\partial \tilde\theta_k]}
 
-    properly concatenated in a single PyTree for every set of parameters. In practice,
-    it should return a result roughly equivalent to the following listing:
+    where now those objects have twice the number of elements as the parameters.
+    In practice, it should return a result roughly equivalent to the following listing:
 
     .. code:: python
 


Test patch:
diff --git a/test/driver/test_srt.py b/test/driver/test_srt.py
--- a/test/driver/test_srt.py
+++ b/test/driver/test_srt.py
@@ -63,7 +63,7 @@ def __call__(self, x):
             return jnp.sum(x, axis=-1).astype(jnp.complex128)
 
 
-def _setup(complex=True):
+def _setup(*, complex=True, machine=None):
     L = 4
     Ns = L * L
     lattice = nk.graph.Square(L, max_neighbor_order=2)
@@ -73,7 +73,8 @@ def _setup(complex=True):
     )
 
     # Define a variational state
-    machine = RBM(num_hidden=2 * Ns, complex=complex)
+    if machine is None:
+        machine = RBM(num_hidden=2 * Ns, complex=complex)
 
     sampler = nk.sampler.MetropolisExchange(
         hilbert=hi,
@@ -94,6 +95,44 @@ def _setup(complex=True):
     return H, opt, vstate
 
 
+def test_SRt_vs_linear_solver_complexpars():
+    """
+    nk.driver.VMC_kernelSR must give **exactly** the same dynamics as nk.driver.VMC with nk.optimizer.SR
+    """
+    n_iters = 5
+
+    model = nk.models.RBM(
+        param_dtype=jnp.complex128,
+        kernel_init=jax.nn.initializers.normal(stddev=0.02),
+        hidden_bias_init=jax.nn.initializers.normal(stddev=0.02),
+        use_visible_bias=False,
+    )
+
+    H, opt, vstate_srt = _setup(machine=model)
+    gs = VMC_SRt(
+        H, opt, variational_state=vstate_srt, diag_shift=0.1, jacobian_mode="complex"
+    )
+    logger_srt = nk.logging.RuntimeLog()
+    gs.run(n_iter=n_iters, out=logger_srt)
+
+    H, opt, vstate_sr = _setup(machine=model)
+    sr = nk.optimizer.SR(solver=solve, diag_shift=0.1, holomorphic=False)
+    gs = nk.driver.VMC(H, opt, variational_state=vstate_sr, preconditioner=sr)
+    logger_sr = nk.logging.RuntimeLog()
+    gs.run(n_iter=n_iters, out=logger_sr)
+
+    # check same parameters
+    jax.tree_map(
+        np.testing.assert_allclose, vstate_srt.parameters, vstate_sr.parameters
+    )
+
+    if mpi.rank == 0:
+        energy_kernelSR = logger_srt.data["Energy"]["Mean"]
+        energy_SR = logger_sr.data["Energy"]["Mean"]
+
+        np.testing.assert_allclose(energy_kernelSR, energy_SR, atol=1e-10)
+
+
 def test_SRt_vs_linear_solver():
     """
     nk.driver.VMC_kernelSR must give **exactly** the same dynamics as nk.driver.VMC with nk.optimizer.SR