nnx and jax.experimental.attrs

[ASEM000]

Hello @cgarciae ,

I am closely watching the progress of nnx, and I wonder what do you think about the new experimental attrs module from jax.

AFAIK you can do something like this with the attrs module

from jax.experimental.attrs import jax_getattr, jax_setattr, register
import jax.numpy as jnp
class Foo:
    def __init__(self):
        self.a = jnp.array([1,2])
        self.b = jnp.array([3,4])

register(Foo)
foo = Foo()

@jax.jit
def swap_ref(foo: Foo):
    a, b = jax_getattr(foo, "a"), jax_getattr(foo, "b")
    jax_setattr(foo, "a", b), jax_setattr(foo, "b", a)

print(foo.a, foo.b) # [1 2] [3 4]
swap_ref(foo)
print(foo.a, foo.b) # [3 4] [1 2]

I know you can achieve the same functionality under nnx.jit, but since jax.jit seems to (yet) to support this feature, do you think using attrs for nnx.Module __setattr__ and __getattr__ will be useful?

[cgarciae]

Hey @ASEM000!

I had a local branch testing attrs support, its very interesting but its not clear that its ultimately the best approach as it has certain limitations: you cannot create new attrs (which we kinda need for sow). I'm very excited for NNX's new graph update mechanism, it gives me way more confidence on recommending the use of nnx.jit as it now properly tracks all state updates, and its now a superset of jax.jit.

[ASEM000]

Thanks for the clarification