Update tutorial examples to thread explicit RNGs

[samanklesaria]

Closes #4971

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[samanklesaria]

When using Jupyter to edit the notebooks, I find that all the {code-cell} blocks in the markdown get replaced with {code-cell} ipython3. That is, jupytext on my machine labels all code cells with the language, while on the main branch, code cells are inconsistently labeled. In the mnist tutorial there's no ipython3 label, but in the randomness tutorial there is. Instead of figuring out how to configure my jupyter settings separately for each notebook file, I'm just going to leave this commit including the extra ipython3 labels.

[cgarciae]

Thanks @samanklesaria ! Its looking great. Left some comments.
Maybe the other place we can update is the index.rst which also contains an example that uses dropout.

[cgarciae]

Thanks @samanklesaria ! Wondering if we could also refactor the intro to the Randomness guide as part of this PR? Current intro is a bit outdated in the sense that it tries to explain it in terms of Haiku and Linen. It might be more useful add a simple code example at the beginning as motivation for the whole guide e.g:

class Model(nnx.Module):
  def __init__(self, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(20, 10, rngs=rngs)
    self.drop = nnx.Dropout(0.1)

  def __call__(self, x, *, rngs):
    return nnx.relu(self.drop(self.linear(x), rngs=rngs))
 
rngs = nnx.Rngs(0)
model = Model(rngs=rngs)  # pass rngs to initialize parameters
x = rngs.normal((32, 20))  # convenient jax.random methods
y = model(x, rngs=rngs)  # pass rngs for dropout masks

We can remove the current section below in favor of the example above and some comments introducing nnx.Rngs as the main mechanism that NNX uses for propagating random state.

Random state handling in Flax NNX was radically simplified compared to systems like Haiku and Flax Linen because Flax NNX defines the random state as an object state. In essence, this means that in Flax NNX, the random state is: 1) just another type of state; 2) stored in nnx.Variables; and 3) held by the models themselves.

The Flax NNX pseudorandom number generator (PRNG) system has the following main characteristics:

It is explicit.

It is order-based.

It uses dynamic counters.

This is a bit different from Flax Linen’s PRNG system, which is (path + order)-based, and uses static counters.

[cgarciae]

Thanks @samanklesaria ! This is great.
As a follow up PR we could change the Fork section to use .split instead.

[Tsukimarf]

ns module methods under jax.named_scope for profiling