[nnx] support Array leaves in graph nodes

[cgarciae]

What does this PR do?

Adds support for Array leaves in graph nodes.

Example:

class Linear(nnx.Module):
  def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
    self.w = jax.random.normal(rngs.params(), (din, dout))
    self.b = jnp.zeros((dout,))

  def __call__(self, x):
    return x @ self.w + self.b