[nnx] add Pytree

[cgarciae]

What does this PR do?

Adds the nnx.Pytree type that can be used to interact more easily with certain JAX APIs that expect callable stateless pytrees such as odeint or even other pytree-based Module systems. nnx.Pytree is a frozen view of any NNX object that is both a pytree and acts as a proxy over the object's methods.

from flax import nnx
from jax.experimental import ode

model = nnx.Linear(2, 2, rngs=nnx.Rngs(0)),
jax_model = nnx.Pytree(model)

y0 = jax.numpy.array([0.5, -0.5])

def dy_dt(y, t, jax_model):
  return jax_model(y)

t = jax.numpy.linspace(0, 1, 10)
y = ode.odeint(dy_dt, y0, t, jax_model)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB