Repository: netket/netket

Issue:
Add `chunk_size` parameter to QGT constructors
This PR adds a `chunk_size` parameter to all three QGT constructors. These allow the QGTs to use a different chunk size than the underlying vstate, which is useful for some architectures, where backpropagation costs more memory than the forward pass. No existing behaviour changes, if `chunk_size` is not supplied, the default is still `vstate.chunk_size`


Discussion:
Nice.

Can you add a test?
Not sure what I should test, all the effects are confined to within the constructor...
# [Codecov](https://codecov.io/gh/netket/netket/pull/1347?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket) Report
> Merging [#1347](https://codecov.io/gh/netket/netket/pull/1347?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket) (a504b49) into [master](https://codecov.io/gh/netket/netket/commit/074ec7c89c9b34310511e71ec4f7a33ce1ce3e18?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket) (074ec7c) will **increase** coverage by `1.35%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #1347      +/-   ##
==========================================
+ Coverage   81.46%   82.82%   +1.35%     
==========================================
  Files         212      212              
  Lines       12638    12637       -1     
  Branches     1927     1928       +1     
==========================================
+ Hits        10296    10467     +171     
+ Misses       1867     1668     -199     
- Partials      475      502      +27     
```


| [Impacted Files](https://codecov.io/gh/netket/netket/pull/1347?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket) | Coverage Δ | |
|---|---|---|
| [netket/optimizer/qgt/qgt\_jacobian\_dense.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L29wdGltaXplci9xZ3QvcWd0X2phY29iaWFuX2RlbnNlLnB5) | `93.40% <100.00%> (+2.10%)` | :arrow_up: |
| [netket/optimizer/qgt/qgt\_jacobian\_pytree.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L29wdGltaXplci9xZ3QvcWd0X2phY29iaWFuX3B5dHJlZS5weQ==) | `96.77% <100.00%> (-0.04%)` | :arrow_down: |
| [netket/optimizer/qgt/qgt\_onthefly.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L29wdGltaXplci9xZ3QvcWd0X29udGhlZmx5LnB5) | `84.14% <100.00%> (+0.19%)` | :arrow_up: |
| [...tket/experimental/hilbert/spin\_orbital\_fermions.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L2V4cGVyaW1lbnRhbC9oaWxiZXJ0L3NwaW5fb3JiaXRhbF9mZXJtaW9ucy5weQ==) | `85.52% <0.00%> (+1.31%)` | :arrow_up: |
| [netket/nn/utils.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L25uL3V0aWxzLnB5) | `86.27% <0.00%> (+1.96%)` | :arrow_up: |
| [netket/tools/info.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L3Rvb2xzL2luZm8ucHk=) | `92.63% <0.00%> (+2.10%)` | :arrow_up: |
| [netket/utils/deprecation.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L3V0aWxzL2RlcHJlY2F0aW9uLnB5) | `93.22% <0.00%> (+5.08%)` | :arrow_up: |
| [netket/tools/check\_mpi.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L3Rvb2xzL2NoZWNrX21waS5weQ==) | `100.00% <0.00%> (+7.14%)` | :arrow_up: |
| [netket/experimental/operator/\_fermions\_2nd.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L2V4cGVyaW1lbnRhbC9vcGVyYXRvci9fZmVybWlvbnNfMm5kLnB5) | `73.59% <0.00%> (+7.47%)` | :arrow_up: |
| [netket/operator/\_pauli\_strings.py](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket#diff-bmV0a2V0L29wZXJhdG9yL19wYXVsaV9zdHJpbmdzLnB5) | `69.75% <0.00%> (+7.71%)` | :arrow_up: |
| ... and [6 more](https://codecov.io/gh/netket/netket/pull/1347/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=netket)

@attila-i-szabo sorry, the weekend made me forget.

Even just a test that checks that the constructor works correctly is fine. I want it around so that we don't break it in the future.

Resolving patch:
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
@@ -69,6 +69,7 @@ def QGTJacobianDense(
             mode=mode,
             holomorphic=holomorphic,
             rescale_shift=rescale_shift,
+            chunk_size=chunk_size,
             **kwargs,
         )
 
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
@@ -70,6 +70,7 @@ def QGTJacobianPyTree(
             mode=mode,
             holomorphic=holomorphic,
             rescale_shift=rescale_shift,
+            chunk_size=chunk_size,
             **kwargs,
         )
 
diff --git a/netket/optimizer/qgt/qgt_onthefly.py b/netket/optimizer/qgt/qgt_onthefly.py
--- a/netket/optimizer/qgt/qgt_onthefly.py
+++ b/netket/optimizer/qgt/qgt_onthefly.py
@@ -47,7 +47,7 @@ def QGTOnTheFly(vstate=None, *, chunk_size=None, **kwargs) -> "QGTOnTheFlyT":
                     memory than the forward pass).
     """
     if vstate is None:
-        return partial(QGTOnTheFly, **kwargs)
+        return partial(QGTOnTheFly, chunk_size=chunk_size, **kwargs)
 
     if "centered" in kwargs:
         warn_deprecation(


Test patch:
diff --git a/test/optimizer/test_sr_api.py b/test/optimizer/test_sr_api.py
--- a/test/optimizer/test_sr_api.py
+++ b/test/optimizer/test_sr_api.py
@@ -14,6 +14,8 @@
 
 import pytest
 
+from typing import Callable
+
 import netket as nk
 
 from .. import common
@@ -37,3 +39,28 @@ def test_deprecated_sr():
 
     with pytest.warns(FutureWarning):
         nk.optimizer.sr.SRLazyGMRES()
+
+
+@pytest.mark.parametrize(
+    "qgt", [nk.optimizer.qgt.QGTJacobianDense, nk.optimizer.qgt.QGTJacobianPyTree]
+)
+def test_qgt_partial_jacobian(qgt):
+    # mode and holomorphic can't be both specified for an actual QGT
+    # but we'll never create one
+    args = {
+        "mode": "real",
+        "holomorphic": False,
+        "diag_shift": 0.03,
+        "rescale_shift": True,
+        "chunk_size": 16,
+    }
+    QGT = qgt(**args)
+    assert isinstance(QGT, Callable)
+    assert QGT.keywords == args
+
+
+def test_qgt_partial_onthefly():
+    args = {"diag_shift": 0.03, "chunk_size": 16}
+    QGT = nk.optimizer.qgt.QGTOnTheFly(**args)
+    assert isinstance(QGT, Callable)
+    assert QGT.keywords == args