Why NNX? understanding the shift in design

[gerdm]

Hi, and thank you for the incredible work on Flax and NNX!

It feels to me that NNX adopts an API and design philosophy reminiscent of PyTorch — for example, removing direct access to weights and moving away from some of the functional programming principles that JAX and Flax, especially Flax’s linen, have embraced. While this shift might make the library more approachable for some, I’m curious about the motivation behind this change.

How does this new approach align with the core principles that made JAX and Flax unique? I noticed that BatchNorm is highlighted in the documentation as a motivating example for this shift, but this raises a question: is the need to simplify specific cases like BatchNorm strong enough to justify a broader change to the API?

Thank you for all the work you’ve put into these tools and for driving innovation in the community!