Util to let bridge module work with NNX submodules

[IvyZX]

• Add current_module to track the current parent
• Add wrap_nnx_module to automate RNG distribution into pure-NNX territory, and to avoid buffer allocation

[cgarciae]

I think we can do a bit better and try to mimic compact with auto-naming or optional name parameter like this:

Suggested change

	layer = nnx.bridge.wrap_nnx_module(lambda r: NNXLayer(8, 0.3, rngs=r))
	setattr(self, f'layer{i}', layer)
	layer = nnx.bridge.compact_init(NNXLayer, 8, 0.3, rngs=self.scope.rngs, name=f'layer{i}')

[cgarciae]

This will break Modules that have keys and counts with additional dimensions e.g.

class FooStack(nnx.Module):
  def __init__(self, n_layers, key):
    keys = jax.random.split(key, n_layers)
    self.rngs = nnx.Rngs(keys)

Looks extreme but users could actually construct something equivalent via vmap.

[cgarciae]

Maybe the policy could be that we only rng state from apply and try to replace it here if the keys and counts are scalars.

[IvyZX]

This is indeed tricky. We supply non-abstract RNG state here essentially for any code in the nnx.Module that assumes a complete nnx.Rngs instance and call its methods in call time.
I think the cleanest solution is to use nnx.jit to partially init the NNX module and only keep all RngState. Updated the PR and added a test for it.