Linen Module attributes and arrays/non hashable

[PhilipVinc]

A thing i noticed is that linen modules can be compared (==) and hashed.
This is useful to me because if I define the function

@jax.jit(static_argnums=0)
def mapply(module, parameters, x):
   return module.apply(parameters, x)

And this will not be re-jitted if I call it sub-sequently with different modules with identical parameters.

However, I now need to add an array as a field to the module (a matrix representing some symmetrical invariances to a dense layer).

class Dense(nn.Module):
   features: int
   symmetrizer: jnp.ndarray
   dtype 

   @compact
   def __call__(self, x)
      x = jnp.asarray(inputs, x)
      kernel = self.param(
            "kernel", self.kernel_init, (inputs.shape[-1], self.features), self.dtype
        )
      kernel = jnp.asarray(kernel, dtype)
      symm_kernel = jnp.einsum("ijkl,kl->ij"self.symmetrizer, kernel)
      x = lax.dot_general(
            x,
            kernel,
            (((inputs.ndim - 1,), (0,)), ((), ())),
            precision=self.precision,
        )
     return x

But the code above will now fail, because module is no longer hashable if it contains an attribute that is a jax array.

This is, I believe because modules declare as 'static' (struct.field(pytree_node=False)) every attribute. Am i right?

What I want to do is build some 'DenseSymm' layer that multiplies the kernel and bias by a symmetrization tensor/matrix.
The Symmetrization tensor will not change and should not be optimised and ideally should be passed in at module construction.

model = DenseSymm(features=10, symmetries=build_symmetrizer_tensor(symmetries)) # or actually RBMSymm, but the api will be the same

I thought it would be a good fit for being an attribute but apparently not... Is this an use case for model_state?

[jheek]

You always make me wonder if I should move back to quantum physics. It all looks exciting but I have no clue what you are doing :P

Modules are not registered as a pytree because they are not simple stateless containers of JAX values.
I think indeed it would be better to register the symmetries as a variable. For example:

class SymmDense(nn.Module):
   features: int
   symm_init: Callable
   dtype

   @compact
   def __call__(self, x)
      x = jnp.asarray(inputs, x)
      kernel = self.param(
            "kernel", self.kernel_init, (inputs.shape[-1], self.features), self.dtype
        )
      kernel = jnp.asarray(kernel, dtype)
      symm_shape = kernel.shape * 2 # I guess?
      symmetrizer = self.variable('constant', 'symmetrizer', self.symm_init, symm_shape)  # assuming symm_init has a shape argument but it does not need to have one...
      symm_kernel = jnp.einsum("ijkl,kl->ij", symmetrizer.value, kernel)
      x = lax.dot_general(
            x,
            kernel,
            (((inputs.ndim - 1,), (0,)), ((), ())),
            precision=self.precision,
        )
     return x

This way you could also enable more advanced things in the feature like using vmap to either map or broadcast over symmetries.

[PhilipVinc]

cc @femtobit for netket#555

[PhilipVinc]

Oh, so we have to cheat and declare it as a simmetrizer? that's... unfortunate.
That also means that one cannot probe the symmetrizer from the model, because that's an initializer.
I don't really like that, but I guess that's the only way?

--

I had started writing the question starting from the physical side, but then thought that it would have been too confusing...
Anyhow, since you asked i'll just paste a condensed version here:

What we are trying to do is re-implement in jax something we had in numpy before, which are symmetric networks.
The wavefunction ψ:ℋ→𝒞 of a quantum system is often invariant under some symmetry transformations T, that takes elements of our domain, for example bitstrings σ=(↑↑↓...), and gives a new one.

The simplest case is translations, where T₁(↑↑↓)=↓↑↑ and T₂(↑↑↓)=↑↓↑, for example, but we don't want to be limited to this simple case.

We use neural networks to approximate the wavefunction, so ψ(σ) is an evaluation of the neural network for a bit-string σ=(↑↑↓...).
We want ψ(σ) = ψ(Tᵢσ) for every i, so essentially ψ(↑↑↓)=ψ(↑↓↑)=ψ(↓↑↑). There are publications showing that this gives better results.

For the permutation case above, this is effectively a konvolution with kernel size as large as the input itself, but again, we don't work only with square lattices which map nicely to convolutions/FFT, and convolutions with large kernels are not efficient (we should do them with FFTs)...
That's why we want the symmetrizer.

[jheek]

Oh, so we have to cheat and declare it as a simmetrizer? that's... unfortunate.
That also means that one cannot probe the symmetrizer from the model, because that's an initializer.
I don't really like that, but I guess that's the only way?

I'm not quite sure what you mean. You can access the value simply as symmetrizer.value which will give you the value as a plain jax array. If you want to access teh symmetrizer outside of call you should define it in setup() so you can access it outside of the Module.__call__ (from other methods or in a parent module for example).

[PhilipVinc]

ah, right. that's setup.
Thank you.