Support for Flax NNX API

Hi Brax development team,

We are currently using the `brax.training.ppo` training scripts in our project and have been very happy with Brax's performance.

We noticed that the `brax.training` module uses the older `flax.linen` API for neural network definitions. Flax recently introduced the `flax.nnx` API, which offers a more Pythonic and streamlined approach to model development.

Are there any plans to transition the brax repository to use the `nnx` API in the future? This would help improve model flexibility and maintainability for projects that depend on `brax.training`.

Thank you for your time and consideration.

Scott