支持 flax.nnx 中的层

我们是否可以在 `einops.layers.flax` 中支持 `flax.nnx` 接口? 当前的 flax 层是基于 `flax.linen` 实现的, flax 已经演进为新的接口: `from flax import nnx`, 请参阅 https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html。 flax 还提供了有关如何将 `flax.linen` 转换为 `flax.nnx` 的指南: https://flax.readthedocs.io/en/latest/guides/bridge_guide.html。