[nnx] add call

[cgarciae]

What does this PR do?

Docs: preview.

Adds nnx.call which takes a (GraphDef, State) pair and creates a proxy object that can be
used to call methods on the underlying graph node. When a method is called, the
output is returned along with a new (GraphDef, State) pair that represents the
updated state of the graph node.

from flax import nnx
import jax
import jax.numpy as jnp

class StatefulLinear(nnx.Module):
  def __init__(self, din, dout, rngs):
    self.w = nnx.Param(jax.random.uniform(rngs(), (din, dout)))
    self.b = nnx.Param(jnp.zeros((dout,)))
    self.count = nnx.Variable(jnp.array(0, dtype=jnp.uint32))

  def __call__(self, x):
    self.count.value += 1
    return x @ self.w + self.b

linear = StatefulLinear(3, 2, nnx.Rngs(0))
linear_state = nnx.split(linear)

@jax.jit
def forward(x, linear_state):
  y, linear_state = nnx.call(linear_state)(x)
  return y, linear_state

x = jnp.ones((1, 3))
y, linear_state = forward(x, linear_state)
y, linear_state = forward(x, linear_state)

linear = nnx.merge(*linear_state)
assert linear.count.value == 2

[PhilipVinc]

Hey @cgarciae ! I'd love this to be merged because I'd like to play with it... Is there any reason this is being left behind?

[cgarciae]

Hey @PhilipVinc, I was wondering if there was a way to avoid increasing the API surface. If you check the updated notes this new version, Pure is now a NamedTuple and nnx.split now returns Pure instances when no filters are passed.