Support layers in flax.nnx

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.