Proper use of nnx.Rngs mixed with random number usage

[jlperla]

I don't yet understand how to mix general random number generation with the NNX Rngs. Can you provide a good pattern to replace the repeated calls to key, subkey = random.split(key, num=2) etc. that people often do with JAX?

To be more concrete, lets say I have

  rngs = nnx.Rngs(0)
  x = jax.random.normal(rngs.next(), (2,))
  model = nnx.Linear(2, 1, rngs=rngs)
  model(x)

  @jax.jit
  def f(key, x):
      return x + jax.random.normal(key, x.shape)
  y = f(rngs.next(), 4)

Note that this uses the rngs for the nnx, but also draws random numbers and calls functions which are expecting normal jax keys.

Is the above correct? is rngs.next() the way to have this reproducible, instead of splitting manually, etc.?