This presentation outlines the upcoming evolution of NNX, the new API for the Flax neural network library, aimed at making it more deeply integrated and "native" to the JAX ecosystem. The speaker addresses current challenges with NNX, where custom transform wrappers and an object-oriented design create a learning curve and friction with JAX's standard PyTree-based paradigm.
The core solution presented is the transition of NNX modules to become native PyTrees. This fundamental change will enable users to apply raw JAX transformations like jax.jit and jax.grad directly, fostering a convergence of the JAX ecosystem around a common standard. This evolution is powered by a new mutable array API in JAX, which allows for more ergonomic, imperative-style code while maintaining JAX's functional purity under the hood. The new design will be more explicit, reducing the framework's complexity and giving users a clearer mental model. These changes will be introduced experimentally behind a feature flag to ensure a smooth transition.
The Evolution of NNX: Becoming More JAX-Native
0:00:05 Current State of NNX: NNX is the new, "Pythonic" API for the Flax library, designed to be intuitive for users coming from frameworks like PyTorch and Keras. However, its custom wrappers and object-oriented nature have created some challenges.
0:01:04 Key Challenges with NNX: The primary issues are a slight learning curve due to transform wrappers that sit on top of JAX, a high maintenance burden for developers, and compatibility friction with JAX's native PyTree-based ecosystem.
0:02:35 Solution Part 1: Adopting PyTrees: In the near future, NNX modules will be refactored to be PyTrees. This makes NNX models natively compatible with the entire JAX ecosystem.
0:02:55 Direct Use of JAX Transforms: A key benefit of this change is that users will be able to apply raw JAX functions like jax.jit and jax.grad directly to NNX modules, eliminating the need for special NNX-specific wrappers.
0:03:18 Ecosystem Convergence: This move is seen as a major step toward unifying the JAX neural network ecosystem, establishing PyTrees as the "lingua franca" and simplifying the landscape compared to older, more complex frameworks.
0:04:23 A More Explicit and Simple API: The new NNX will be more explicit. Instead of the framework automatically handling state, users will explicitly split differentiable and non-differentiable parts. This reduces the framework's complexity and gives users a clearer mental model.
0:05:34 Solution Part 2: The New Mutable Array API: This shift is enabled by a new mutable array API in JAX. It allows developers to write imperative-looking code (e.g., mutating a value) that can be JIT-compiled, providing ergonomic Pythonic syntax while compiling to pure, functional code.
0:07:01 Gradual Rollout Plan: The new functionality will be introduced experimentally behind a mutable_array feature flag. It will not be the default initially, allowing time for documentation and guides to be prepared for a smooth transition.
0:07:48 Key Takeaway: NNX is becoming more JAX-native. By adopting PyTrees and leveraging JAX's new mutable arrays, the framework is simplifying its API, reducing its surface area, and making the user experience much closer to writing raw JAX code.