Design Decision on the nnx library

[ymahlau]

Hi everyone,

I would like to start a discussion about the design descisions of the nnx library. Specifically, I asked myself why reference sharing and mutability was introduced as this kind of goes against the design of JAX itself.

The only reason I can see to try to support mutability and reference sharing within JAX is to enable a more "pytorch-like" interface. And NNX has definitely suceeded with this. In my opinion, the implementing models with nnx feels very smooth and avoids a lot of the annoyences of flax.linen.

However, the problems start when trying to integrate nnx modules with more complicated jax transformations and existing libraries. This can also be seen by many of the current issues in this repo.

As another point, mutability and reference sharing are not strictly necessary to hold the state inside the model itself, as demonstrated by the serket library.

Therefore my question would be what others think the advantages and disadvantages of mutability and reference sharing are. Also, maybe some developers could give more information about the design decisions they had to face.

[cgarciae]

Hey @ymahlau, thanks for bringing this up!

However, the problems start when trying to integrate nnx modules with more complicated jax transformations and existing libraries. This can also be seen by many of the current issues in this repo.

Agreed, we've been thinking about this and have a plan that will allow NNX to be used with raw JAX transformations and other libraries by making NNX objects implement the pytree protocol. Will publish an announcement soon-ish.

[Artoriuz]

It appears to me that NNX is getting progressively closer to Equinox, in a way that they're becoming somewhat redundant.

I'm saying this because we've seen a few other JAX ML libraries come and go (Trax, Haiku, etc, etc...), and it looks like we already have what NNX is trying to achieve somewhere else.

Am I missing something? Is there anything in NNX's design that clearly sets it apart from Equinox?

[cgarciae]

They have different philosophies, NNX is focused on supporting python reference semantics while Equinox prefers immutability and value semantics.

[carlosgmartin]

Related: #4431.

[adam-hartshorne]

Any update on this new functionality to allow NNX to work with raw JAX transformations and other libraries?