Longterm plans for NNX

[thomaspinder]

Hi!

I'm curious what the long-term plan is for NNX? Is the roadmap for it to eventually replace the use of struct? Or will it coexist?

Thanks!

[chiamp]

@cgarciae

[cgarciae]

Hey @thomaspinder ! It's still very early to tell how NNX will evolve, as stated in the README Linen will continue to be the main API. NNX is recommended for newer / more experimental projects or people who enjoy the benefits of its API.

Regarding struct, they are not very related, you can use NNX pytrees like State inside a struct.dataclass instance. nnx is more of a replacement for the linen module.

[thomaspinder]

Thanks @cgarciae. I ask as I develop GPJax and we're looking to build on top of Flax moving forward. I would love to use NNX, however, the risk of the module being deprecated concerns me.

[cgarciae]

@thomaspinder I wouldn't worry about deprecation, our intention is to have NNX or something similar be the main API at some point. However, there is a long road ahead.

[AllanChain]

@cgarciae May I ask does this mean that the linen API will be deprecated eventually? I'm concerned because the shape inference feature from linen API is quite handy, and there is no replacement in the NNX API.

[francois-rozet]

Hello @cgarciae, I see that nnx came a long way since we last exchanged. Great job! If you remember, I have been working on Inox to simplify the API of Equinox and I now believe the interface is almost ready. Interestingly, there are many similarities with nnx, but also key differences. Notably,

1. In inox, modules are PyTrees, which means they can be passed in and out of native JAX transformations. They can also be serialized (pickle) out of the box.
2. Like in nnx, internal arrays have roles (Parameter, Statistic, ...), indicated by a wrapper. However, I designed the interface such that a Variable instance is a valid JAX array (jax.core.valid_jaxtype) and can used directly with jax.numpy (and jax.lax). Consequently, modules don't need custom __getattribute__ and __setattr__ (which are prone to silent errors) and variables can be stored in arbitrarily complex PyTrees (e.g. a list of dicts of tuples).

If you are interested, you can check inox.nn.module and inox.tree_util in the documentation. Tell me what you think!

[cgarciae]

Hey @francois-rozet! Thanks for reaching out.

I like it a lot! Its arguably a better version of Treex with some very similar ideas to NNX. NNX also has a somewhat hidden pytree abstraction which looks very similar to Inox Module: pytreelib.py.

The key difference is that NNX Modules don't implement the Pytree protocol (its trivial, there is an experimental implementation) so they can support mutable shared references and its underlaying abstraction is a Graph instead of a Tree. Pytree libraries have a very hard time representing mutable shared state due to referential transparency. Check out our Why NNX? doc, I'd love to hear your feedback.

[francois-rozet]

I actually have a work-in-progress branch to support shared references:

https://github.com/francois-rozet/inox/blob/share/inox/nn/sharing.py

The idea is to force the user to explicitly indicate which objects are shared references within a scope (a special module) and to prune all but one reference when flattening. From my (very short) tests, it works for graphs as well. I have a harder time with repr of graphs actually 😅

Check out our Why NNX? doc, I'd love to hear your feedback.

I already did it, but will look more in depth! BTW, what is the print(f'{x = }) syntax? Seems very useful!

[cgarciae]

Did you see this section of Why aren't Modules Pytrees? Our main problem is that with Pytrees shared references get lost in the process. I also tried a custom flattening to preserve shared references a while back in jax-ml/jax#7919, but I found I couldn't support is_leaf for pytree operations when I was doing the flattening myself instead of JAX.

BTW: why don't all Modules behave like Scope? And maybe even Reference.

[francois-rozet]

Indeed, if the tree_flatten method flattens the whole tree instead of a single level, you cannot use the is_leaf and with_path interfaces, which is very annoying. This is why Scope prunes the tree with tree_map to return the same as a standard module. However, this means that if all modules were self-pruning, it would be quite expensive (at least outside of jax.jit).

Also there is still a problem if all shared references are not flattened at the same time. However, I think this should be easy to detect (and warn the user).

[carlosgmartin]

@cgarciae Does Flax still accept PRs for Linen, or only for NNX from now on? (I ask because I have some open PRs pertaining to Linen that haven't been merged after several months.)