Jax from flax.nn to flax.linen?

[dave-andersen]

Hi, folks - thanks for the useful release!

flax.nn is now deprecated, so a move to linen would help ensure this code remains future compatible (and doesn't generate warnings, as it does now). Not sure if you're intending to maintain this or leave it as a snapshot of what the paper was based upon, but it's a nice resource to have out there.