Support layers in flax.nnx

[gongzhitaao]

Can we support flax.nnx interface in einops.layers.flax?

Current flax layer is implemented based upon flax.linen, The flax has evolved into a new interface: from flax import nnx, see https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html. Flax also provides guides on how to convert flax.linen to flax.nnx: https://flax.readthedocs.io/en/latest/guides/bridge_guide.html.

[arogozhnikov]

Well-tested contribution are welcome, @cgarciae probably can help with reviews. Also, if you like jax/flax and want to work on it - contact Cristian.