Repository: netket/netket

Issue:
Passing masks to a GCNN has no effect
Hi, first of all I'm sorry if this actually belongs in the discussions section, I've never worked with github before and don't know what problems are Issues-worthy.

I noticed that after the new implementation of the masked GCNNs, passing masks to the GCNN does not have any effect, e.g. for a model with only 2nd nearest neighbors convolutions the number of parameters does not change, so
```
# Define system
L = 8
lattice = nk.graph.Square(length=L)
hi = nk.hilbert.Spin(s=1 / 2, total_sz=0, N=lattice.n_nodes)

# Define Metropolis-Hastings sampler
sampler = nk.sampler.MetropolisExchange(hilbert=hi, graph=lattice)

# Define the model: masked GCNN
input_mask = np.zeros([L, L])
for i in range(-1, 2):
    for j in range(-1, 2):
        input_mask[i][j] = 1
input_mask = input_mask.ravel()
hidden_mask = np.repeat(np.expand_dims(input_mask, 1), repeats=8, axis=1).ravel()

machine = nk.models.GCNN(symmetries=lattice, layers=2, features=(4, 2), param_dtype=jnp.complex128,
                         input_mask=input_mask, hidden_mask=hidden_mask)
vstate = nk.vqs.MCState(sampler=sampler, model=machine)
print("number of parameters:", vstate.n_parameters)
```
returns `number of parameters: 4358` instead of the desired `number of parameters: 618`


I think the error is that in `netket.models.equivariant` , the masks are not passed on from the general GCNN method  to the constructors of the different modes (FFT, irreps).

Also, if I use the DenseSymm layer with a mask I get an error, so changing the model in the above code to 
`machine = nk.nn.DenseSymm(symmetries=lattice, mode="matrix", mask=HashableArray(input_mask), features=2)
` 
gives the error `TypeError: nonzero requires ndarray or scalar arguments, got <class 'netket.utils.array.HashableArray'> at position 0.`

It should work if `self.kernel_indices = jnp.nonzero(self.mask)[0]`  in the `DenseSymmMatrix` class
is changed to `(self.kernel_indices,) = np.nonzero(self.mask)`, as in the other modes.






Discussion:
Very sorry for the late response, but thanks so much for identifying these bugs. I will write a PR to fix them in the near future. If you have any other questions feel free to @ me.

Resolving patch:
diff --git a/netket/models/equivariant.py b/netket/models/equivariant.py
--- a/netket/models/equivariant.py
+++ b/netket/models/equivariant.py
@@ -789,6 +789,8 @@ def GCNN(
                 parity=parity,
                 param_dtype=param_dtype,
                 complex_output=complex_output,
+                hidden_mask=hidden_mask,
+                input_mask=input_mask,
                 **kwargs,
             )
         else:
@@ -801,6 +803,8 @@ def GCNN(
                 shape=shape,
                 param_dtype=param_dtype,
                 complex_output=complex_output,
+                hidden_mask=hidden_mask,
+                input_mask=input_mask,
                 **kwargs,
             )
     elif mode in ["irreps", "auto"]:
@@ -819,6 +823,8 @@ def GCNN(
                 parity=parity,
                 param_dtype=param_dtype,
                 complex_output=complex_output,
+                hidden_mask=hidden_mask,
+                input_mask=input_mask,
                 **kwargs,
             )
         else:
@@ -830,6 +836,8 @@ def GCNN(
                 characters=characters,
                 param_dtype=param_dtype,
                 complex_output=complex_output,
+                hidden_mask=hidden_mask,
+                input_mask=input_mask,
                 **kwargs,
             )
     else:
diff --git a/netket/nn/symmetric_linear.py b/netket/nn/symmetric_linear.py
--- a/netket/nn/symmetric_linear.py
+++ b/netket/nn/symmetric_linear.py
@@ -74,7 +74,7 @@ def setup(self):
         # pylint: disable=attribute-defined-outside-init
         self.n_symm, self.n_sites = np.asarray(self.symmetries).shape
         if self.mask is not None:
-            self.kernel_indices = jnp.nonzero(self.mask)[0]
+            (self.kernel_indices,) = np.nonzero(self.mask.wrapped)
 
     @compact
     def __call__(self, x: Array) -> Array:
@@ -181,7 +181,7 @@ def setup(self):
         self.sites_per_cell = sg.shape[1] // self.n_cells
 
         if self.mask is not None:
-            (self.kernel_indices,) = np.nonzero(self.mask)
+            (self.kernel_indices,) = np.nonzero(self.mask.wrapped)
 
         # maps (n_sites) dimension of kernels to (sites_per_cell, n_point, *shape)
         # as used in FFT-based group convolution
@@ -312,7 +312,7 @@ def setup(self):
         self.n_cells = np.product(np.asarray(self.shape))
         self.n_point = len(pt) // self.n_cells
         if self.mask is not None:
-            (self.kernel_indices,) = np.nonzero(self.mask)
+            (self.kernel_indices,) = np.nonzero(self.mask.wrapped)
 
         # maps (n_sites) dimension of kernels to (n_point, n_point, *shape)
         # as used in FFT-based group convolution
@@ -447,7 +447,7 @@ class DenseEquivariantIrrep(Module):
     def setup(self):
         self.n_symm = self.irreps[0].shape[0]
         if self.mask is not None:
-            (self.kernel_indices,) = np.nonzero(self.mask)
+            (self.kernel_indices,) = np.nonzero(self.mask.wrapped)
 
         self.forward = jnp.concatenate(
             [jnp.asarray(irrep).reshape(self.n_symm, -1) for irrep in self.irreps],
@@ -611,7 +611,7 @@ class DenseEquivariantMatrix(Module):
     def setup(self):
         self.n_symm = np.asarray(self.product_table).shape[0]
         if self.mask is not None:
-            (self.kernel_indices,) = np.nonzero(self.mask)
+            (self.kernel_indices,) = np.nonzero(self.mask.wrapped)
 
     @compact
     def __call__(self, x: Array) -> Array:
@@ -674,7 +674,9 @@ def __call__(self, x: Array) -> Array:
 
 
 @deprecate_dtype
-def DenseSymm(symmetries, point_group=None, mode="auto", shape=None, **kwargs):
+def DenseSymm(
+    symmetries, point_group=None, mode="auto", shape=None, mask=None, **kwargs
+):
     r"""
     Implements a projection onto a symmetry group. The output will be
     equivariant with respect to the symmetry operations in the group and can
@@ -713,6 +715,9 @@ def DenseSymm(symmetries, point_group=None, mode="auto", shape=None, **kwargs):
         bias_init: Optional bias initialization function. Defaults to zero initialization.
 
     """
+    if mask is not None:
+        mask = HashableArray(mask)
+
     if isinstance(symmetries, Lattice) and (
         point_group is not None or symmetries._point_group is not None
     ):
@@ -740,9 +745,9 @@ def DenseSymm(symmetries, point_group=None, mode="auto", shape=None, **kwargs):
                 "the symmetries keyword argument."
             )
         else:
-            return DenseSymmFFT(sym, shape=shape, **kwargs)
+            return DenseSymmFFT(sym, shape=shape, mask=mask, **kwargs)
     elif mode in ["matrix", "auto"]:
-        return DenseSymmMatrix(sym, **kwargs)
+        return DenseSymmMatrix(sym, mask=mask, **kwargs)
     else:
         raise ValueError(
             f"Unknown mode={mode}. Valid modes are 'fft', 'matrix', or 'auto'."
@@ -757,6 +762,7 @@ def DenseEquivariant(
     shape=None,
     point_group=None,
     in_features=None,
+    mask=None,
     **kwargs,
 ):
     r"""A group convolution operation that is equivariant over a symmetry group.
@@ -801,6 +807,9 @@ def DenseEquivariant(
         kernel_init: Optional kernel initialization function. Defaults to variance scaling.
         bias_init: Optional bias initialization function. Defaults to zero initialization.
     """
+    if mask is not None:
+        mask = HashableArray(mask)
+
     # deprecate in_features
     if in_features is not None:
         warn_deprecation(
@@ -851,13 +860,13 @@ def DenseEquivariant(
     elif isinstance(symmetries, Sequence):
         if mode not in ["irreps", "auto"]:
             raise ValueError("Specification of symmetries incompatible with mode")
-        return DenseEquivariantIrrep(symmetries, **kwargs)
+        return DenseEquivariantIrrep(symmetries, mask=mask, **kwargs)
     else:
         if symmetries.ndim == 2 and symmetries.shape[0] == symmetries.shape[1]:
             if mode == "irreps":
                 raise ValueError("Specification of symmetries incompatible with mode")
             elif mode == "matrix":
-                return DenseEquivariantMatrix(symmetries, **kwargs)
+                return DenseEquivariantMatrix(symmetries, mask=mask, **kwargs)
             else:
                 if shape is None:
                     raise TypeError(
@@ -866,7 +875,9 @@ def DenseEquivariant(
                         "the symmetries keyword argument."
                     )
                 else:
-                    return DenseEquivariantFFT(symmetries, shape=shape, **kwargs)
+                    return DenseEquivariantFFT(
+                        symmetries, mask=mask, shape=shape, **kwargs
+                    )
         return ValueError("Invalid Specification of Symmetries")
 
     if mode == "fft":
@@ -878,13 +889,15 @@ def DenseEquivariant(
             )
         else:
             return DenseEquivariantFFT(
-                HashableArray(sg.product_table), shape=shape, **kwargs
+                HashableArray(sg.product_table), mask=mask, shape=shape, **kwargs
             )
     elif mode in ["irreps", "auto"]:
         irreps = tuple(HashableArray(irrep) for irrep in sg.irrep_matrices())
-        return DenseEquivariantIrrep(irreps, **kwargs)
+        return DenseEquivariantIrrep(irreps, mask=mask, **kwargs)
     elif mode == "matrix":
-        return DenseEquivariantMatrix(HashableArray(sg.product_table), **kwargs)
+        return DenseEquivariantMatrix(
+            HashableArray(sg.product_table), mask=mask, **kwargs
+        )
     else:
         raise ValueError(
             f"Unknown mode={mode}. Valid modes are 'fft', 'matrix', 'irreps' or 'auto'."
diff --git a/netket/utils/array.py b/netket/utils/array.py
--- a/netket/utils/array.py
+++ b/netket/utils/array.py
@@ -32,10 +32,12 @@ class HashableArray:
     """The wrapped array. Note that this array is read-only."""
 
     def __pre_init__(self, wrapped):
-        wrapped = wrapped.copy()
-
-        if isinstance(wrapped, np.ndarray):
-            wrapped.flags.writeable = False
+        if isinstance(wrapped, HashableArray):
+            wrapped = wrapped.wrapped
+        else:
+            wrapped = wrapped.copy()
+            if isinstance(wrapped, np.ndarray):
+                wrapped.flags.writeable = False
 
         return (wrapped,), {}
 


Test patch:
diff --git a/test/models/test_gcnn.py b/test/models/test_gcnn.py
--- a/test/models/test_gcnn.py
+++ b/test/models/test_gcnn.py
@@ -94,6 +94,52 @@ def test_gcnn(mode, complex_output):
     vmc.advance(1)
 
 
+def n_parameters(ma):
+    params = ma.init(jax.random.PRNGKey(0), np.zeros(4))
+
+
+def test_gcnn_mask():
+
+    g = nk.graph.Chain(4)
+
+    hi = nk.hilbert.Spin(s=1 / 2, total_sz=0, N=g.n_nodes)
+    sa = nk.sampler.MetropolisExchange(hilbert=hi, graph=g)
+
+    input_mask = np.concatenate((np.zeros(2), np.ones(2)), 0)
+    hidden_mask = np.concatenate((np.zeros(4), np.ones(4)), 0)
+
+    ma1 = nk.models.GCNN(
+        symmetries=g,
+        layers=2,
+        features=2,
+        input_mask=None,
+        hidden_mask=None,
+    )
+
+    ma2 = nk.models.GCNN(
+        symmetries=g,
+        layers=2,
+        features=2,
+        input_mask=input_mask,
+        hidden_mask=None,
+    )
+
+    ma3 = nk.models.GCNN(
+        symmetries=g,
+        layers=2,
+        features=2,
+        input_mask=None,
+        hidden_mask=hidden_mask,
+    )
+
+    vstate1 = nk.vqs.MCState(sampler=sa, model=ma1)
+    vstate2 = nk.vqs.MCState(sampler=sa, model=ma2)
+    vstate3 = nk.vqs.MCState(sampler=sa, model=ma3)
+
+    assert vstate1.n_parameters - vstate2.n_parameters == 4
+    assert vstate1.n_parameters - vstate3.n_parameters == 16
+
+
 @pytest.mark.parametrize("mode", ["fft", "irreps"])
 def test_GCNN_creation(mode):
 
diff --git a/test/models/test_nn.py b/test/models/test_nn.py
--- a/test/models/test_nn.py
+++ b/test/models/test_nn.py
@@ -40,6 +40,58 @@ def _setup_symm(symmetries, N, lattice=nk.graph.Chain):
     return g, hi, perms
 
 
+@pytest.mark.parametrize("mode", ["fft", "matrix"])
+def test_DenseSymm_matrixInput(mode):
+    g, hi, perms = _setup_symm("trans", N=8)
+    shape = tuple(g.extent)
+    perms_matrix = np.array(perms)
+    ma = nk.nn.DenseSymm(
+        symmetries=perms,
+        mode=mode,
+        features=2,
+        shape=shape,
+    )
+    ma_matrix = nk.nn.DenseSymm(
+        symmetries=perms_matrix,
+        mode=mode,
+        features=2,
+        shape=shape,
+    )
+
+    assert hash(ma) == hash(ma_matrix)
+
+    k = jax.random.PRNGKey(0)
+    x = hi.numbers_to_states(np.array([0, 1, 2])).reshape(3, 1, hi.size)
+    p1 = jax.jit(ma.init)(k, x)
+    p2 = jax.jit(ma_matrix.init)(k, x)
+    assert nk.jax.tree_size(p1) == nk.jax.tree_size(p2)
+
+    n_symm = perms_matrix.shape[0]
+    mask = np.zeros(n_symm)
+    mask[np.random.choice(n_symm, n_symm // 2, replace=False)] = 1
+
+    ma_masked = nk.nn.DenseSymm(
+        symmetries=perms,
+        mode=mode,
+        features=2,
+        shape=shape,
+        mask=mask,
+    )
+    p3 = jax.jit(ma_masked.init)(k, x)
+    assert nk.jax.tree_size(p3) < nk.jax.tree_size(p1)
+
+    ma_masked2 = nk.nn.DenseSymm(
+        symmetries=perms,
+        mode=mode,
+        features=2,
+        shape=shape,
+        mask=nk.utils.HashableArray(mask),
+    )
+    p4 = jax.jit(ma_masked2.init)(k, x)
+    assert nk.jax.tree_size(p4) == nk.jax.tree_size(p3)
+    assert hash(ma_masked) == hash(ma_masked2)
+
+
 @pytest.mark.parametrize("symmetries", ["trans", "space_group"])
 @pytest.mark.parametrize("use_bias", [True, False])
 @pytest.mark.parametrize("mode", ["fft", "matrix"])
diff --git a/test/utils/test_utils.py b/test/utils/test_utils.py
--- a/test/utils/test_utils.py
+++ b/test/utils/test_utils.py
@@ -63,6 +63,12 @@ def test_HashableArray(numpy):
     assert_equal(wa.wrapped, np.asarray(wa))
     assert wa.wrapped is not wa
 
+    # test construction from hashable array
+    wa2 = HashableArray(a)
+    assert hash(wa) == hash(wa2)
+    assert wa == wa2
+    assert_equal(wa2.wrapped, np.asarray(wa))
+
 
 def test_Kahan_sum():
     ksum1 = KahanSum(0.0)