Update JAX example to use Flax NNX API

[SauravMaheshkar]

Describe the type of feature and its functionality.

The current recommended way to use Flax is to via the modern NNX API.

Describe step by step what files and adjustments are you planning to include.

This can either be done by updating examples/quickstart-jax or by creating a new example.

Is there something else you want to add?

Happy to take this up

[WilliamLindskog]

Hi @SauravMaheshkar,

great suggestion! Do you see a reason to keep the quickstart-jax example as is for now and create another one (similar to the current one but based on NNX API)? What do you think?

[SauravMaheshkar]

IMO there isn't much reason to keep the legacy API example, most modern JAX implementations use the NNX API.

[WilliamLindskog]

@SauravMaheshkar Want to open a PR to update the example? Happy to support this

[SauravMaheshkar]

Will put together a PR