[P] Introducing NNX: Neural Networks for JAX

Can we have the power of Flax with the simplicity of Equinox?

[NNX](https://github.com/cgarciae/nnx) is a highly experimental 🧪 proof of concept framework that provides Pytree Modules with:

* Shared state
* Tractable mutability
* Semantic partitioning (collections)

Defining Modules is very similar to Equinox, but you mark parameters with `nnx.param`, this creates some Refx references under the hood. Similar to flax, you use `make_rng` to request RNG keys which you seed during `init`.


[Linear Module](https://preview.redd.it/2kf6ff5aahwa1.png?width=1506&format=png&auto=webp&v=enabled&s=d3a942676fd5171bc9a2481d589abf602d6365bb)

NNX introduces the concept of …

concept example jax machinelearning modules mutability networks neural networks partitioning power semantic simplicity state transformation