internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: pytorch/tensordict

Issue:
[BUG] Wrong batch_size
## Describe the bug

The shape of a batch is the full length of the dataset

## To Reproduce

I am using version `=0.0.3`

```python
import torch
from tensordict import MemmapTensor
from tensordict.prototype import tensorclass
from torch.utils.data import DataLoader

@tensorclass
class Data:
    images: torch.Tensor
    targets: torch.Tensor

    @classmethod
    def from_tensors(cls,  device: torch.device):
        len_dataset = 1000
        images, targets = torch.randn((len_dataset, 28, 28)), torch.randn((len_dataset, 1))
        data = cls(
            images=MemmapTensor.from_tensor(images, filename="images.dat"),
            targets=MemmapTensor.from_tensor(targets, filename="targets.dat"),
            batch_size=[len_dataset],
            device=device,
        )

        data.memmap_()

        return data
    
data = Data.from_tensors(device="cuda")

dl = DataLoader(data, batch_size=64, collate_fn=lambda x: x)

for batch in dl:
    print(batch.images.shape)
```

The output is, should be `64` as first shape (the batch_size)

```bash
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
torch.Size([1000, 28, 28])
```



Discussion:

Resolving patch:
diff --git a/tensordict/prototype/tensorclass.py b/tensordict/prototype/tensorclass.py
--- a/tensordict/prototype/tensorclass.py
+++ b/tensordict/prototype/tensorclass.py
@@ -173,6 +173,7 @@ def __torch_function__(
     cls.__setattr__ = _setattr_wrapper(cls.__setattr__, expected_keys)
     cls.__getattr__ = _getattr
     cls.__getitem__ = _getitem
+    cls.__getitems__ = _getitem
     cls.__setitem__ = _setitem
     cls.__repr__ = _repr
     cls.__len__ = _len


Test patch:
Repository: pytorch/tensordict

Issue:
Default batch_size + TensorDictModule can be surprising
TensordictModule allows to use a tensordict/class as input to the forward module. But it also support the nn.Module style call with tensors.
Under the hood it creates a new [tensordict](https://github.com/pytorch-labs/tensordict/blob/main/tensordict/nn/common.py#L253) that then forwards to the call method. In doing so, however, it uses the [default batch_size](https://github.com/pytorch-labs/tensordict/blob/e92f26366472a47aff93a0838716af4f5349191f/tensordict/tensordict.py#L3012).

I would say that using the maximum size is probably not the most common case. Can we reconsider the default or find a way to get the right size from the model itself?


Discussion:
I guess your point is that then writing onto that tensordict can be problematic if the output shape violates it, right?
Sure I think we can remove that
exactly

Resolving patch:
diff --git a/tensordict/nn/common.py b/tensordict/nn/common.py
--- a/tensordict/nn/common.py
+++ b/tensordict/nn/common.py
@@ -250,7 +250,9 @@ def wrapper(_self, *args: Any, **kwargs: Any) -> Any:
                                 f"The key {expected_key} wasn't found in the keyword arguments "
                                 f"but is expected to execute that function."
                             )
-                tensordict = make_tensordict(tensordict_values)
+                tensordict = make_tensordict(
+                    tensordict_values, batch_size=torch.Size([])
+                )
                 out = func(_self, tensordict, *args, **kwargs)
                 out = tuple(out[key] for key in dest)
                 return out[0] if len(out) == 1 else out


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -520,6 +520,14 @@ def test_dispatch(self):
         out = tdm(a=torch.zeros(1, 1))
         assert (out == td["b"]).all()
 
+    def test_dispatch_changing_size(self):
+        # regression test on non max batch-size for dispatch
+        tdm = TensorDictModule(nn.Linear(1, 2), ["a"], ["b"])
+        td = TensorDict({"a": torch.zeros(1, 1)}, 1)
+        tdm(td)
+        out = tdm(a=torch.zeros(1, 1))
+        assert (out == td["b"]).all()
+
     def test_dispatch_nested(self):
         tdm = TensorDictModule(nn.Linear(1, 1), [("a", "c")], [("b", "d")])
         td = TensorDict({("a", "c"): torch.zeros(1, 1)}, [1])
Repository: netket/netket

Issue:
Get rid of netket::Qubits class
`netket::Qubits` is really just a `CustomHilbert` with `localstates = {0.0, 1.0}`. No need for code duplication.


Discussion:
Agreed, this should be implemented in pure python actually 

Resolving patch:
diff --git a/netket/hilbert.py b/netket/hilbert.py
--- a/netket/hilbert.py
+++ b/netket/hilbert.py
@@ -1 +1,22 @@
 from ._C_netket.hilbert import *
+
+
+def Qubit(graph):
+    """
+    Constructs a new ``Qubit`` given a graph.
+
+    Args:
+        graph: Graph representation of sites.
+
+    Examples:
+        Simple qubit hilbert space.
+
+        ```python
+        >>> from netket.graph import Hypercube
+        >>> from netket.hilbert import Qubit
+        >>> g = Hypercube(length=10,n_dim=2,pbc=True)
+        >>> hi = Qubit(graph=g)
+        >>> print(hi.size)
+        100
+    """
+    return CustomHilbert(graph, local_states=[0, 1])


Test patch:
Repository: netket/netket

Issue:
module 'netket' has no attribute 'MPI'
With the merge #193 we have lost the MPI module 


Discussion:

Resolving patch:
diff --git a/netket/__init__.py b/netket/__init__.py
--- a/netket/__init__.py
+++ b/netket/__init__.py
@@ -31,3 +31,4 @@
     utils,
     variational,
 )
+from ._C_netket import MPI, LookupReal, LookupComplex


Test patch:
Repository: netket/netket

Issue:
Getting acceptance of numpy metropolis_hastings
The method defining the acceptance for `numpy/metropolis_hastings` in v3.0 is

```python
@property
def acceptance(self):
    """The measured acceptance probability."""
    return _mean(self._accepted_samples) / _mean(self._total_samples)
```

But `self._accepted_samples` and `self._total_samples` are ints. Don't know if this has been fixed in a PR, but if it's not, maybe it should be included.


Discussion:
willing to fix this in a PR? thanks!
Sure thing. I'll do that tomorrow morning.

Resolving patch:
diff --git a/netket/sampler/numpy/metropolis_hastings.py b/netket/sampler/numpy/metropolis_hastings.py
--- a/netket/sampler/numpy/metropolis_hastings.py
+++ b/netket/sampler/numpy/metropolis_hastings.py
@@ -1,7 +1,8 @@
 from ..abstract_sampler import AbstractSampler
-from ...stats import mean as _mean
 from netket import random as _random
 
+from netket.stats import sum_inplace as _sum_inplace
+
 import math
 import numpy as _np
 from numba import jit, int64, float64
@@ -127,4 +128,4 @@ def __next__(self):
     @property
     def acceptance(self):
         """The measured acceptance probability."""
-        return _mean(self._accepted_samples) / _mean(self._total_samples)
+        return _sum_inplace(self._accepted_samples) / _sum_inplace(self._total_samples)
diff --git a/netket/stats/_sum_inplace.py b/netket/stats/_sum_inplace.py
--- a/netket/stats/_sum_inplace.py
+++ b/netket/stats/_sum_inplace.py
@@ -31,6 +31,7 @@ def sum_inplace(x):
 @sum_inplace.register(_np.complex64)
 @sum_inplace.register(_np.complex128)
 @sum_inplace.register(float)
+@sum_inplace.register(int)
 def sum_inplace_scalar(a):
     ar = _np.asarray(a)
 


Test patch:
Repository: netket/netket

Issue:
ExactSampler is not resetting at construction time
Exact Sampler should call Reset at construction time 



Discussion:

Resolving patch:
diff --git a/netket/sampler/exact_sampler.py b/netket/sampler/exact_sampler.py
--- a/netket/sampler/exact_sampler.py
+++ b/netket/sampler/exact_sampler.py
@@ -33,6 +33,7 @@ def __init__(self, machine, sample_size=16):
         else:
             self.hilbert = machine.hilbert
         self._machine_pow = 2.0
+        self.reset()
 
     def reset(self, init_random=False):
         self._prob = _np.absolute(self.machine.to_array()) ** self.machine_pow


Test patch:
Repository: netket/netket

Issue:
Floating Point Error in Lattice.py
I tried generating a periodic Kagome lattice using lattice.py:

>>> kagome = nk.graph.Lattice(basis_vectors=[[2.,0.],[1.,np.sqrt(3)]],extent=[2,2],atoms_coord=[[0.,0.],[1./2.,np.sqrt(3)/2.],[1.,0.]])

Only half of the edges appeared:

>>>kagome.n_edges()
12

The bug is related to floating points, some distances are registered as 0.99999... and some are registered as 1.0



Discussion:

Resolving patch:
diff --git a/netket/graph/lattice.py b/netket/graph/lattice.py
--- a/netket/graph/lattice.py
+++ b/netket/graph/lattice.py
@@ -32,11 +32,13 @@ def get_edges(atoms_positions, cutoff):
         min_dists[node] = min_dist
     for node in _np.unique(id1):
         min_dist = _np.min(values[id1 == node])
-        mask = (id1 == node) & (values == min_dist)
+        mask = (id1 == node) & (_np.isclose(values, min_dist))
         first = id1[mask]
         second = id2[mask]
         for pair in zip(first, second):
-            if min_dist == min_dists[pair[0]] and min_dist == min_dists[pair[1]]:
+            if _np.isclose(min_dist, min_dists[pair[0]]) and _np.isclose(
+                min_dist, min_dists[pair[1]]
+            ):
                 pairs.append(pair)
     return pairs
 


Test patch:
diff --git a/Test/GroundState/test_vmc.py b/Test/GroundState/test_vmc.py
--- a/Test/GroundState/test_vmc.py
+++ b/Test/GroundState/test_vmc.py
@@ -37,7 +37,7 @@ def _setup_vmc(dtype=np.float32, sr=True):
         sr_config = nk.optimizer.SR()
     else:
         sr_config = None
-    driver = nk.Vmc(ha, op, variational_state=vs, sr=sr_config)
+    driver = nk.VMC(ha, op, variational_state=vs, sr=sr_config)
 
     return ha, sx, vs, sa, driver
 
@@ -106,7 +106,7 @@ def test_vmc_construction_vstate():
 
     with raises(TypeError):
         ha2 = nk.operator.LocalOperator(ha.hilbert * ha.hilbert)
-        driver = nk.Vmc(ha2, op, variational_state=driver.state)
+        driver = nk.VMC(ha2, op, variational_state=driver.state)
 
 
 def test_vmc_functions():
Repository: netket/netket

Issue:
Duplicate edges and self-loops in `Lattice` with PBC
The `Lattice` class and, by consequence, the new `Grid` implementation, adds self-loops and duplicate edges when PBC are enabled:

Consider the following cases:
```python
>>> print(Grid([1]).edges())
[(0, 0)]
>>> print(Grid([2, 1]).edges())
[(0, 0), (0, 1), (0, 1), (1, 1)]
>>> print(Grid([3, 2]).edges())
[(0, 2), (0, 4), (0, 1), (0, 1), (1, 3), (1, 5), (2, 3), (2, 3), (2, 4), (3, 5), (4, 5), (4, 5)]
```

While one can argue that the first two cases are ill-formed and should just throw and error, the third case does make sense if not for the fact that additional periodic edges are added along the short direction.

Since PBC along that direction do not make sense, but OTOH it would be tedious to require users to manually specify `pbc=[True, False]` in such as case in my opinion, I think this should just default to assuming `pbc=False` along directions of `extent <= 2` (this would also address the first two problems).

I can make a PR for this, if everyone agrees this is the correct solution.


Discussion:

Resolving patch:
diff --git a/netket/graph/lattice.py b/netket/graph/lattice.py
--- a/netket/graph/lattice.py
+++ b/netket/graph/lattice.py
@@ -230,7 +230,9 @@ def __init__(
                 of length `ndim`.
             pbc: If ``True`` then the constructed lattice
                 will have periodic boundary conditions, otherwise
-                open boundary conditions are imposed (default=`True`).
+                open boundary conditions are imposed. Can also be an boolean sequence
+                of length `ndim`, indicating either open or closed boundary conditions
+                separately for each direction.
             site_offsets: The position offsets of sites in the unit cell (one site at
                 the origin by default).
             distance_atol: Distance below which spatial points are considered equal for
@@ -305,6 +307,8 @@ def __init__(
         graph = igraph.Graph()
         graph.add_vertices(len(old_nodes))
         graph.add_edges([(new_nodes[edge[0]], new_nodes[edge[1]]) for edge in edges])
+        graph.simplify()
+
         self._sites = []
         for i, site in enumerate(sites[old_node] for old_node in old_nodes):
             site.id = i


Test patch:
diff --git a/Test/Graph/test_graph.py b/Test/Graph/test_graph.py
--- a/Test/Graph/test_graph.py
+++ b/Test/Graph/test_graph.py
@@ -120,6 +120,22 @@ def test_lattice_graphs(i, name):
     assert graph.n_edges == graph.n_nodes * coordination_number[i] // 2
 
 
+# netket#743 : multiple edges and self-loops in lattice
+@pytest.mark.parametrize(
+    "size,n_nodes,n_edges",
+    zip(
+        [[1], [2, 1], [3, 2]],
+        [1, 2, 6],
+        [0, 1, 9],
+    ),
+)
+def test_no_redundant_edges(size, n_nodes, n_edges):
+    g = Grid(size)
+    print(g.edges())
+    assert g.n_nodes == n_nodes
+    assert g.n_edges == n_edges
+
+
 @pytest.mark.parametrize("g", graphs + symmetric_graphs)
 def test_lattice(g):
     if not isinstance(g, Lattice):
Repository: netket/netket

Issue:
Docstring for ```sampler.is_exact``` is not... exact 
[This docstring](https://github.com/netket/netket/blob/ad832f07470dc30bd299ec2ee833860f62b1c63d/netket/sampler/base.py#L151) is a bit ambiguous: it says 

> The sampler is exact if the samples are always distributed according to the        chosen power of the variational state. 

However I would argue that the sampler is exact if all the samples are exactly distributed according to the chosen power of the variational state and (also, crucially) **there is no correlation among them**


Discussion:
Ok, let me change this

Resolving patch:
diff --git a/netket/sampler/base.py b/netket/sampler/base.py
--- a/netket/sampler/base.py
+++ b/netket/sampler/base.py
@@ -152,8 +152,8 @@ def is_exact(self) -> bool:
         """
         Returns `True` if the sampler is exact.
 
-        The sampler is exact if the samples are always distributed according to the
-        chosen power of the variational state.
+        The sampler is exact if all the samples are exactly distributed according to the
+        chosen power of the variational state, and there is no correlation among them.
         """
         return False
 


Test patch:
Repository: netket/netket

Issue:
AbstractVariationalDriver.py bug in updating step_count
![image](https://user-images.githubusercontent.com/47376937/125520970-8a189955-4de6-45b2-acfc-45eab4c82d9e.png)

I think this error can be fixed either by writing a setter function for the `@property step_count` in the same way that `@optimizer.setter` exists in the same file. 

Otherwise, maybe this can be fixed by changing ` self.step_count = 0` to ` self._step_count = 0`.

![image](https://user-images.githubusercontent.com/47376937/125521480-ca41de24-c283-4d8c-bf16-ba25a93b2399.png)



Discussion:

Resolving patch:
diff --git a/netket/driver/abstract_variational_driver.py b/netket/driver/abstract_variational_driver.py
--- a/netket/driver/abstract_variational_driver.py
+++ b/netket/driver/abstract_variational_driver.py
@@ -112,7 +112,7 @@ def reset(self):
         count is set to 0.
         """
         self.state.reset()
-        self.step_count = 0
+        self._step_count = 0
         pass
 
     @abc.abstractmethod


Test patch:
Repository: netket/netket

Issue:
Docs for sigma operators
The documentation of the ladder operator `netket.operator.spin.sigmap` says `σ^+ = σ^x + i σ^y`, but in fact it is defined as `σ^+ = 1/2(σ^x + i σ^y)` in the code.

The same error is in the documentation of `netket.operator.spin.sigmam` operator.


Discussion:
@mmezic Thanks for spotting this, you are right. This is also confirmed by the tests:
https://github.com/netket/netket/blob/00d4607806885f0320ab6ddead9bf497aea268c5/test/operator/test_spin.py#L24-L27

The docs should be fixed.

Resolving patch:
diff --git a/netket/operator/spin.py b/netket/operator/spin.py
--- a/netket/operator/spin.py
+++ b/netket/operator/spin.py
@@ -88,7 +88,7 @@ def sigmaz(hilbert: AbstractHilbert, site: int, dtype: DType = float) -> _LocalO
 
 def sigmam(hilbert: AbstractHilbert, site: int, dtype: DType = float) -> _LocalOperator:
     """
-    Builds the :math:`\\sigma^{-} = \\sigma^x - i \\sigma^y` operator acting on the
+    Builds the :math:`\\sigma^{-} = \\frac{1}{2}(\\sigma^x - i \\sigma^y)` operator acting on the
     `site`-th of the Hilbert space `hilbert`.
 
     If `hilbert` is a non-Spin space of local dimension M, it is considered
@@ -111,7 +111,7 @@ def sigmam(hilbert: AbstractHilbert, site: int, dtype: DType = float) -> _LocalO
 
 def sigmap(hilbert: AbstractHilbert, site: int, dtype: DType = float) -> _LocalOperator:
     """
-    Builds the :math:`\\sigma^{+} = \\sigma^x + i \\sigma^y` operator acting on the
+    Builds the :math:`\\sigma^{+} = \\frac{1}{2}(\\sigma^x + i \\sigma^y)` operator acting on the
     `site`-th of the Hilbert space `hilbert`.
 
     If `hilbert` is a non-Spin space of local dimension M, it is considered


Test patch:
diff --git a/test/operator/test_spin.py b/test/operator/test_spin.py
--- a/test/operator/test_spin.py
+++ b/test/operator/test_spin.py
@@ -1,7 +1,8 @@
+import pytest
+
 from netket.operator import spin
 import netket as nk
 import numpy as np
-
 from numpy.testing import assert_almost_equal
 
 herm_operators = {}
@@ -10,8 +11,9 @@
 N = 4
 
 
-def test_pauli_algebra():
-    hi = nk.hilbert.Spin(1 / 2) ** N
+@pytest.mark.parametrize("S", [1 / 2, 1, 3 / 2])
+def test_pauli_algebra(S):
+    hi = nk.hilbert.Spin(S) ** N
 
     for i in range(N):
         sx = spin.sigmax(hi, i)
@@ -26,8 +28,11 @@ def test_pauli_algebra():
         assert_almost_equal(0.5 * (sx.to_dense() - 1j * sy.to_dense()), sm.to_dense())
         assert_almost_equal(0.5 * (sx.to_dense() + 1j * sy.to_dense()), sp.to_dense())
 
-        Imat = np.eye(hi.n_states)
+        if S == 1 / 2:
+            Imat = np.eye(hi.n_states)
 
-        # check that -i sx sy sz = I
-        assert_almost_equal((-1j * sx @ sy @ sz).to_dense(), Imat)
-        assert_almost_equal((-1j * sx.to_dense() @ sy.to_dense() @ sz.to_dense()), Imat)
+            # check that -i sx sy sz = I
+            assert_almost_equal((-1j * sx @ sy @ sz).to_dense(), Imat)
+            assert_almost_equal(
+                (-1j * sx.to_dense() @ sy.to_dense() @ sz.to_dense()), Imat
+            )
Repository: netket/netket

Issue:
Broken equations in docs
See for example the unredered tex in https://www.netket.org/docs/sr.html 


Discussion:

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -62,11 +62,10 @@
 source_suffix = [".rst", ".ipynb", ".md"]
 
 # Markdown parser latex support
-myst_enable_extensions = ["dollarmath", "amsmath"]
-
+myst_enable_extensions = ["dollarmath", "amsmath", "braket"]
+myst_update_mathjax = False
 mathjax_path = "https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"
 
-
 # -- Options for HTML output -------------------------------------------------
 
 html_theme = "bootstrap"


Test patch:
Repository: netket/netket

Issue:
[input validation] nk.graph.Chain([L]) does not fial
nk.graph.Chain([L]) should fail, but does not, and leads to failure later that it took me a while to realise...

We should validate the input to be an integer.

```python
>>> graph = nk.graph.Chain([L])
>>> sgb = graph.space_group_builder()
>>> sgb.little_group([0])
```


Discussion:

Resolving patch:
diff --git a/netket/graph/common_lattices.py b/netket/graph/common_lattices.py
--- a/netket/graph/common_lattices.py
+++ b/netket/graph/common_lattices.py
@@ -140,6 +140,8 @@ def Hypercube(length: int, n_dim: int = 1, *, pbc: bool = True) -> Lattice:
          >>> print(g.n_nodes)
          1000
     """
+    if not isinstance(length, int) or length <= 0:
+        raise TypeError("Argument `length` must be a positive integer")
     length_vector = [length] * n_dim
     return Grid(length_vector, pbc=pbc)
 


Test patch:
diff --git a/test/graph/test_graph.py b/test/graph/test_graph.py
--- a/test/graph/test_graph.py
+++ b/test/graph/test_graph.py
@@ -305,6 +305,18 @@ def test_graph_wrong():
     with pytest.raises(ValueError):
         nk.graph.Graph([1, 2, 3], [1, 2, 3])
 
+    with pytest.raises(TypeError):
+        nk.graph.Hypercube([5])
+
+    with pytest.raises(TypeError):
+        nk.graph.Cube([5])
+
+    with pytest.raises(TypeError):
+        nk.graph.Square([5])
+
+    with pytest.raises(TypeError):
+        nk.graph.Chain([5])
+
 
 def test_edges_are_correct():
     def check_edges(length, n_dim, pbc):
Repository: netket/netket

Issue:
`get_conn` should throw an error on wrong input size
Currently get_conn does no input shape validation and then reshapes the input carelessly.
So if you give it a batch of states, which is not supported, he gives back garbage becase

```
def get_conn
        return self.get_conn_flattened(
            x.reshape((1, -1)),
            np.ones(1),
        )
```


Discussion:

Resolving patch:
diff --git a/netket/operator/_abstract_operator.py b/netket/operator/_abstract_operator.py
--- a/netket/operator/_abstract_operator.py
+++ b/netket/operator/_abstract_operator.py
@@ -186,7 +186,7 @@ def get_conn_flattened(
         """
         raise NotImplementedError()
 
-    def get_conn(self, x):
+    def get_conn(self, x: np.ndarray):
         r"""Finds the connected elements of the Operator. Starting
         from a given quantum number x, it finds all other quantum numbers x' such
         that the matrix element :math:`O(x,x')` is different from zero. In general there
@@ -200,7 +200,17 @@ def get_conn(self, x):
             matrix: The connected states x' of shape (N_connected,hilbert.size)
             array: An array containing the matrix elements :math:`O(x,x')` associated to each x'.
 
+        Raise:
+            ValueError: If the given quantum number is not compatible with the hilbert space.
         """
+        if x.ndim != 1:
+            raise ValueError(
+                "get_conn does not support batches. Please use get_conn_flattened instead."
+            )
+        if x.shape[0] != self.hilbert.size:
+            raise ValueError(
+                "The given quantum numbers do not match the hilbert space."
+            )
 
         return self.get_conn_flattened(
             x.reshape((1, -1)),


Test patch: