Repository: netket/netket

Issue:
nkjax.jacobian has an undocumented rescaling when center=True
See https://github.com/netket/netket/blob/45d334e94eb4088550cc697eac19525919bed607/netket/jax/_jacobian/logic.py#L314

When `center=True` we rescale by `sqrt(n_samples)` after centering.

This is done consistently when a pdf is passed in, but only when centering if no pdf is passed in.

As this function is called `jacobian`, I think the correct behaviour would be to not divide by this rescaling factor, and embed it in higher level users such as QGT & co.


Discussion:

Resolving patch:
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
@@ -32,7 +32,15 @@
 
 
 @partial(
-    jax.jit, static_argnames=("apply_fun", "mode", "chunk_size", "center", "dense")
+    jax.jit,
+    static_argnames=(
+        "apply_fun",
+        "mode",
+        "chunk_size",
+        "center",
+        "dense",
+        "_sqrt_rescale",
+    ),
 )
 def jacobian(
     apply_fun: Callable,
@@ -45,6 +53,7 @@ def jacobian(
     chunk_size: int = None,
     center: bool = False,
     dense: bool = False,
+    _sqrt_rescale: bool = False,
 ) -> PyTree:
     r"""
     Computes the jacobian of a NN model with respect to its parameters. This function
@@ -121,6 +130,13 @@ def jacobian(
             are the derivatives wrt the real part of the parameters, while the
             second :math:`N_\text{pars}` elements are the derivatives wrt the
             imaginary part of the paramters.
+        _sqrt_rescale: **internal flag** (do not rely on it) a boolean flag
+            (disabled by default). If enabled, the jacobian is rescaled by
+            :math:`1/\sqrt{N_\text{samples}}` to match the scaling emerging in
+            some use-cases such when building the Quantum Geometric Tensor.
+            If a pdf is specified, the scaling will instead be
+            :math:`\sqrt{pdf_i}`. This flag is temporary and internal and might
+            be discontinued at any point in the future. Do not use it.
 
 
     Extra details of the different modes are given below:
@@ -308,11 +324,14 @@ def jacobian(
     )(f, params, samples)
 
     if pdf is None:
-        sqrt_n_samp = math.sqrt(samples.shape[0] * mpi.n_nodes)  # maintain weak type
         if center:
-            jacobians = jax.tree_map(
-                lambda x: subtract_mean(x, axis=0) / sqrt_n_samp, jacobians
-            )
+            jacobians = jax.tree_map(lambda x: subtract_mean(x, axis=0), jacobians)
+
+        if _sqrt_rescale:
+            sqrt_n_samp = math.sqrt(
+                samples.shape[0] * mpi.n_nodes
+            )  # maintain weak type
+            jacobians = jax.tree_map(lambda x: x / sqrt_n_samp, jacobians)
 
     else:
         if center:
@@ -321,7 +340,8 @@ def jacobian(
             )
             jacobians = jax.tree_map(lambda x, y: x - y, jacobians, jacobians_avg)
 
-        jacobians = _multiply_by_pdf(jacobians, jnp.sqrt(pdf))
+        if _sqrt_rescale:
+            jacobians = _multiply_by_pdf(jacobians, jnp.sqrt(pdf))
 
     return jacobians
 
diff --git a/netket/optimizer/qgt/qgt_jacobian_dense.py b/netket/optimizer/qgt/qgt_jacobian_dense.py
--- a/netket/optimizer/qgt/qgt_jacobian_dense.py
+++ b/netket/optimizer/qgt/qgt_jacobian_dense.py
@@ -137,6 +137,7 @@ def QGTJacobianDense(
         chunk_size=chunk_size,
         dense=True,
         center=True,
+        _sqrt_rescale=True,
     )
 
     if offset is not None:
diff --git a/netket/optimizer/qgt/qgt_jacobian_pytree.py b/netket/optimizer/qgt/qgt_jacobian_pytree.py
--- a/netket/optimizer/qgt/qgt_jacobian_pytree.py
+++ b/netket/optimizer/qgt/qgt_jacobian_pytree.py
@@ -139,6 +139,7 @@ def QGTJacobianPyTree(
         chunk_size=chunk_size,
         dense=False,
         center=True,
+        _sqrt_rescale=True,
     )
 
     if offset is not None:


Test patch:
diff --git a/test/optimizer/test_qgt_logic.py b/test/optimizer/test_qgt_logic.py
--- a/test/optimizer/test_qgt_logic.py
+++ b/test/optimizer/test_qgt_logic.py
@@ -408,6 +408,7 @@ def apply_fun(params, samples):
         mode=mode,
         dense=False,
         center=True,
+        _sqrt_rescale=True,
     )
     # TODO Apply offset if offset is not None