[nnx] add support for standalone Variables

[cgarciae]

What does this PR do?

Adds support for using Variables outside of graph nodes, meaning they can now be passed directly to transforms, including inside pytrees without a parent graph node, and the ability to call split, merge, and update on them directly.

rngs = nnx.Rngs(0)
w = nnx.Param(jax.random.normal(rngs(), (2, 3)))
b = nnx.Param(jnp.zeros((3,)))
count = nnx.Variable(jnp.array(0))

@nnx.jit
def linear(w, b, count, x):
  count += 1
  return x @ w + b[None]

x = jax.random.normal(rngs(), (1, 2))
y = linear(w, b, count, x)

assert count.value == 1
assert y.shape == (1, 3)

[github-advanced-security]

This pull request sets up GitHub code scanning for this repository. Once the scans have completed and the checks have passed, the analysis results for this pull request branch will appear on this overview. Once you merge this pull request, the 'Security' tab will show more code scanning analysis results (for example, for the default branch). Depending on your configuration and choice of analysis tool, future pull requests will be annotated with code scanning analysis results. For more information about GitHub code scanning, check out the documentation.

[IvyZX]

This lambda becomes kinda long & unreadable. Can you refactor it out to a function? Then you can also use this for the merge_fn call in line 374.

[cgarciae]

great idea! done