Passing masks to a GCNN has no effect

Created by: jobdky

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.