Encapsulate Arrays inside Variables

[cgarciae]

Hey @maxencefaldor! This is a very cool project, I love it.

Just wanted to give you a heads up that in the next release of Flax, all state is required to be inside Variables. This solves some issues regarding attributes changing from being state or not due to the type. As an easy fix you can encapsulate Array inside nnx.Param or create a custom nnx.Variable type e.g.

self.kernel_fft = nnx.Param(jnp.fft.fft2(jnp.fft.fftshift(kernel_normalized, axes=(0, 1)), axes=(0, 1)))

Most of the code should work as Variables behave as proxies over their inner value and can be passed to JAX functions. The main difference is that to update them you must use .value:

self.kernel_fft.value = new_value

Best,
Cristian