Repository: netket/netket

Issue:
`QGTJacobianDense` does not work with models using single precision(float32/complex64) parameters
This was brought to my attention right now.
I think it happens because JacobianDense concatenates the parameters and the way we do this plays badly when the parameters are single precision but the output of the model is double precision...

Should be relatively easy to fix, if someone wanted to pick it up (@inailuig ?)

Note: this came up by someone using a GCNN with single precision parameters. The default QGT implementation in that case picks jacobiandense which then hits this bug.

Reproducer
```python
import netket as nk
import jax.numpy as jnp

# 1D Lattice
L = 20
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)

# Ising spin hamiltonian
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)

# RBM Spin Machine
ma = nk.models.RBM(alpha=1, dtype=jnp.float32)

# Metropolis Local Sampling
sa = nk.sampler.MetropolisLocal(hi, n_chains=16)

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.1)

# SR
sr = nk.optimizer.SR(qgt=nk.optimizer.qgt.QGTJacobianDense(diag_shift=0.01))

# Variational state
vs = nk.vqs.MCState(sa, ma, n_samples=1000, n_discard_per_chain=100)

# Variational monte carlo driver with a variational state
gs = nk.VMC(ha, op, variational_state=vs, preconditioner=sr)

# Run the optimization for 300 iterations
gs.run(n_iter=300, out=None)
```
`QGTJacobianDense` does not work with models using single precision(float32/complex64) parameters
This was brought to my attention right now.
I think it happens because JacobianDense concatenates the parameters and the way we do this plays badly when the parameters are single precision but the output of the model is double precision...

Should be relatively easy to fix, if someone wanted to pick it up (@inailuig ?)

Note: this came up by someone using a GCNN with single precision parameters. The default QGT implementation in that case picks jacobiandense which then hits this bug.

Reproducer
```python
import netket as nk
import jax.numpy as jnp

# 1D Lattice
L = 20
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)

# Ising spin hamiltonian
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)

# RBM Spin Machine
ma = nk.models.RBM(alpha=1, dtype=jnp.float32)

# Metropolis Local Sampling
sa = nk.sampler.MetropolisLocal(hi, n_chains=16)

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.1)

# SR
sr = nk.optimizer.SR(qgt=nk.optimizer.qgt.QGTJacobianDense(diag_shift=0.01))

# Variational state
vs = nk.vqs.MCState(sa, ma, n_samples=1000, n_discard_per_chain=100)

# Variational monte carlo driver with a variational state
gs = nk.VMC(ha, op, variational_state=vs, preconditioner=sr)

# Run the optimization for 300 iterations
gs.run(n_iter=300, out=None)
```


Discussion:
Some more things which don't work as expected in dense and pytree:

- oks are double precision when both params and samples are single precision 
- mixed parameter dtypes are not reconstructed correctly 
But mixed parameters dtypes are supposed not to work
Some more things which don't work as expected in dense and pytree:

- oks are double precision when both params and samples are single precision 
- mixed parameter dtypes are not reconstructed correctly 
But mixed parameters dtypes are supposed not to work

Resolving patch:
diff --git a/netket/optimizer/qgt/qgt_jacobian_dense_logic.py b/netket/optimizer/qgt/qgt_jacobian_dense_logic.py
--- a/netket/optimizer/qgt/qgt_jacobian_dense_logic.py
+++ b/netket/optimizer/qgt/qgt_jacobian_dense_logic.py
@@ -287,7 +287,9 @@ def gradf_fun(params, σ):
     )
 
     n_samp = samples.shape[0] * mpi.n_nodes
-    centered_oks = subtract_mean(jacobians, axis=0) / np.sqrt(n_samp)
+    centered_oks = subtract_mean(jacobians, axis=0) / np.sqrt(
+        n_samp, dtype=jacobians.dtype
+    )
 
     centered_oks = centered_oks.reshape(-1, centered_oks.shape[-1])
 
diff --git a/netket/optimizer/qgt/qgt_jacobian_pytree_logic.py b/netket/optimizer/qgt/qgt_jacobian_pytree_logic.py
--- a/netket/optimizer/qgt/qgt_jacobian_pytree_logic.py
+++ b/netket/optimizer/qgt/qgt_jacobian_pytree_logic.py
@@ -114,7 +114,7 @@ def _divide_by_sqrt_n_samp(oks, samples):
     divide Oⱼₖ by √n
     """
     n_samp = samples.shape[0] * mpi.n_nodes  # MPI
-    return jax.tree_map(lambda x: x / np.sqrt(n_samp), oks)
+    return jax.tree_map(lambda x: x / np.sqrt(n_samp, dtype=x.dtype), oks)
 
 
 def _multiply_by_pdf(oks, pdf):


Test patch:
diff --git a/test/optimizer/test_qgt_itersolve.py b/test/optimizer/test_qgt_itersolve.py
--- a/test/optimizer/test_qgt_itersolve.py
+++ b/test/optimizer/test_qgt_itersolve.py
@@ -56,9 +56,20 @@
 solvers_tol = {}
 
 solvers["gmres"] = partial(jax.scipy.sparse.linalg.gmres, tol=1e-6)
-solvers_tol[solvers["gmres"]] = 4e-4
+solvers_tol[solvers["gmres"], np.dtype("float64")] = 5e-4, 0
+solvers_tol[solvers["gmres"], np.dtype("float32")] = 1e-2, 1e-4
 solvers["cholesky"] = nk.optimizer.solver.cholesky
-solvers_tol[solvers["cholesky"]] = 1e-8
+solvers_tol[solvers["cholesky"], np.dtype("float64")] = 1e-8, 0
+solvers_tol[solvers["cholesky"], np.dtype("float32")] = 1e-2, 1e-4
+
+matmul_tol = {}
+matmul_tol[np.dtype("float64")] = 1e-7, 0
+matmul_tol[np.dtype("float32")] = 5e-4, 0
+
+
+dense_tol = {}
+dense_tol[np.dtype("float64")] = 1e-5, 1e-15
+dense_tol[np.dtype("float32")] = 5e-4, 1e-6
 
 
 RBM = partial(
@@ -73,13 +84,12 @@
 )
 
 models = {
-    "RBM[dtype=float]": partial(RBM, dtype=float),
-    "RBM[dtype=complex]": partial(RBM, dtype=complex),
-    "RBMModPhase[dtype=float]": partial(RBMModPhase, dtype=float),
+    "RBM[dtype=float64]": partial(RBM, dtype=np.dtype("float64")),
+    "RBM[dtype=float32]": partial(RBM, dtype=np.dtype("float32")),
+    "RBM[dtype=complex128]": partial(RBM, dtype=np.dtype("complex128")),
+    "RBMModPhase[dtype=float64]": partial(RBMModPhase, dtype=np.dtype("float64")),
 }
 
-dtypes = {"float": float, "complex": complex}
-
 
 @pytest.fixture(
     params=[pytest.param(modelT, id=name) for name, modelT in models.items()]
@@ -96,11 +106,17 @@ def vstate(request, model, chunk_size):
     N = 5
     hi = nk.hilbert.Spin(1 / 2, N)
 
+    k = jax.random.PRNGKey(3)
+    k1, k2 = jax.random.split(k)
+
     vstate = nk.vqs.MCState(
         nk.sampler.MetropolisLocal(hi),
         model,
+        seed=k1,
     )
-    vstate.init_parameters(normal(stddev=0.001), seed=jax.random.PRNGKey(3))
+
+    # initialize the same parameters on every rank
+    vstate.init_parameters(normal(stddev=0.001), seed=k2)
 
     vstate.sample()
 
@@ -122,8 +138,9 @@ def test_qgt_solve(qgt, vstate, solver, _mpi_size, _mpi_rank):
     S = qgt(vstate)
     x, _ = S.solve(solver, vstate.parameters)
 
+    rtol, atol = solvers_tol[solver, nk.jax.dtype_real(vstate.model.dtype)]
     jax.tree_map(
-        partial(testing.assert_allclose, rtol=solvers_tol[solver]),
+        partial(testing.assert_allclose, rtol=rtol, atol=atol),
         S @ x,
         vstate.parameters,
     )
@@ -143,7 +160,9 @@ def test_qgt_solve(qgt, vstate, solver, _mpi_size, _mpi_rank):
             x_all, _ = S.solve(solver, vstate.parameters)
 
             jax.tree_map(
-                lambda a, b: np.testing.assert_allclose(a, b, rtol=0.00045), x, x_all
+                lambda a, b: np.testing.assert_allclose(a, b, rtol=rtol, atol=atol),
+                x,
+                x_all,
             )
 
 
@@ -167,6 +186,9 @@ def test_qgt_solve_with_x0(qgt, vstate):
 )
 @pytest.mark.parametrize("chunk_size", [None, 16])
 def test_qgt_matmul(qgt, vstate, _mpi_size, _mpi_rank):
+
+    rtol, atol = matmul_tol[nk.jax.dtype_real(vstate.model.dtype)]
+
     S = qgt(vstate)
     rng = nkjax.PRNGSeq(0)
     y = jax.tree_map(
@@ -185,7 +207,11 @@ def check_same_dtype(x, y):
     x_dense = S @ y_dense
     x_dense_unravelled = unravel(x_dense)
 
-    jax.tree_map(lambda a, b: np.testing.assert_allclose(a, b), x, x_dense_unravelled)
+    jax.tree_map(
+        lambda a, b: np.testing.assert_allclose(a, b, rtol=rtol, atol=atol),
+        x,
+        x_dense_unravelled,
+    )
 
     if _mpi_size > 1:
         # other check
@@ -201,7 +227,11 @@ def check_same_dtype(x, y):
             S = qgt(vstate)
             x_all = S @ y
 
-            jax.tree_map(lambda a, b: np.testing.assert_allclose(a, b), x, x_all)
+            jax.tree_map(
+                lambda a, b: np.testing.assert_allclose(a, b, rtol=rtol, atol=atol),
+                x,
+                x_all,
+            )
 
 
 @pytest.mark.parametrize(
@@ -210,6 +240,9 @@ def check_same_dtype(x, y):
 )
 @pytest.mark.parametrize("chunk_size", [None, 16])
 def test_qgt_dense(qgt, vstate, _mpi_size, _mpi_rank):
+
+    rtol, atol = dense_tol[nk.jax.dtype_real(vstate.model.dtype)]
+
     S = qgt(vstate)
 
     # test repr
@@ -242,7 +275,7 @@ def test_qgt_dense(qgt, vstate, _mpi_size, _mpi_rank):
             S = qgt(vstate)
             Sd_all = S.to_dense()
 
-            np.testing.assert_allclose(Sd_all, Sd, rtol=1e-5, atol=1e-15)
+            np.testing.assert_allclose(Sd_all, Sd, rtol=rtol, atol=atol)
 
 
 @pytest.mark.skipif_mpi