[nnx] add tabulate

[cgarciae]

What does this PR do?

Adds nnx.tabulate with a similar API to linen.tabulate.

Example:

class Block(nnx.Module):
  def __init__(self, din, dout, rngs: nnx.Rngs):
    self.linear = nnx.Linear(din, dout, rngs=rngs)
    self.bn = nnx.BatchNorm(dout, rngs=rngs)
    self.dropout = nnx.Dropout(0.2, rngs=rngs)

  def __call__(self, x):
    return nnx.relu(self.dropout(self.bn(self.linear(x))))

class Foo(nnx.Module):
  def __init__(self, rngs: nnx.Rngs):
    self.block1 = Block(32, 128, rngs=rngs)
    self.block2 = Block(128, 10, rngs=rngs)

  def __call__(self, x):
    return self.block2(self.block1(x))

foo = Foo(nnx.Rngs(0))

print(nnx.tabulate(foo, jnp.ones((1, 32))))