Switch network implementations from `Flax.linen` to `Flax.NNX`

[nic-barbara]

This seems to be the future of neural network models in JAX. See migration instructions here: https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html

[nic-barbara]

Having read through this, the switch to Flax.NNX will massively simplify the code. For example:

• We can have non-jittable code in the __init__ which means we won't need a separate setup method to initialise a REN from an explicit model (currently left as future work/in development).
• With nnx, we can call methods of a model directly without needing model.apply(params, ..., method="<insert_method_name>").

At the moment, I've chosen to stick with Flax.linen just because nnx is so new and so few packages support it. However, the documentation is excellent, and we should certainly consider migrating over in the near future.