Change introduction to randomness

Author: samanklesaria

docs_nnx/guides/randomness.ipynb

@@ -6,19 +6,7 @@
    "source": [
     "# Randomness\n",
     "\n",
-    "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.Variable`s; and 3) held by the models themselves.\n",
-    "\n",
-    "The Flax NNX [pseudorandom number generator (PRNG)](https://flax.readthedocs.io/en/latest/glossary.html#term-RNG-sequences) system has the following main characteristics:\n",
-    "\n",
-    "- It is **explicit**.\n",
-    "- It is **order-based**.\n",
-    "- It uses **dynamic counters**.\n",
-    "\n",
-    "This is a bit different from [Flax Linen's PRNG system](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html), which is `(path + order)`-based, and uses static counters.\n",
-    "\n",
-    "> **Note:** To learn more about random number generation in JAX, the `jax.random` API, and PRNG-generated sequences, check out this [JAX PRNG tutorial](https://jax.readthedocs.io/en/latest/random-numbers.html).\n",
-    "\n",
-    "Let’ start with some necessary imports:"
+    "Flax NNX uses the stateful `nnx.Rngs` class to simplify Jax's handling of random states. For example, the code below uses a `nnx.Rngs` object to define a simple linear model with dropout:"
    ]
   },
   {
@@ -28,14 +16,35 @@
    "outputs": [],
    "source": [
     "from flax import nnx\n",
-    "import jax\n",
-    "from jax import random, numpy as jnp"
+    "\n",
+    "class Model(nnx.Module):\n",
+    "  def __init__(self, *, rngs: nnx.Rngs):\n",
+    "    self.linear = nnx.Linear(20, 10, rngs=rngs)\n",
+    "    self.drop = nnx.Dropout(0.1)\n",
+    "\n",
+    "  def __call__(self, x, *, rngs):\n",
+    "    return nnx.relu(self.drop(self.linear(x), rngs=rngs))\n",
+    "\n",
+    "rngs = nnx.Rngs(0)\n",
+    "model = Model(rngs=rngs)  # pass rngs to initialize parameters\n",
+    "x = rngs.normal((32, 20))  # convenient jax.random methods\n",
+    "y = model(x, rngs=rngs)  # pass rngs for dropout masks"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "We always pass `nnx.Rngs` objects to models at initialization (to initialize parameters). For models with nondeterministic outputs like the one above, we also pass `nnx.Rngs` objects to the model's `__call__` method.\n",
+    "\n",
+    "The Flax NNX [pseudorandom number generator (PRNG)](https://flax.readthedocs.io/en/latest/glossary.html#term-RNG-sequences) system has the following main characteristics:\n",
+    "\n",
+    "- It is **explicit**.\n",
+    "- It is **order-based**.\n",
+    "- It uses **dynamic counters**.\n",
+    "\n",
+    "> **Note:** To learn more about random number generation in JAX, the `jax.random` API, and PRNG-generated sequences, check out this [JAX PRNG tutorial](https://jax.readthedocs.io/en/latest/random-numbers.html).\n",
+    "\n",
     "## `Rngs`, `RngStream`, and `RngState`\n",
     "\n",
     "In Flax NNX, the `nnx.Rngs` type is the primary convenience API for managing the random state(s). Following Flax Linen's footsteps, `nnx.Rngs` have the ability to create multiple named PRNG key [streams](https://jax.readthedocs.io/en/latest/jep/263-prng.html), each with its own state, for the purpose of having tight control over randomness in the context of [JAX transformations (transforms)](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations).\n",
@@ -199,6 +208,7 @@
     }
    ],
    "source": [
+    "import jax.numpy as jnp\n",
     "dropout(jnp.ones(4), rngs=rngs)"
    ]
   },
@@ -295,6 +305,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import jax\n",
     "rngs = nnx.Rngs(0, params=1)\n",
     "\n",
     "# using jax.random\n",
@@ -539,7 +550,7 @@
     "\n",
     "In Flax NNX, there are two ways to approach this:\n",
     "\n",
-    "1. By passing an `nnx.Rngs` object through the `__call__` stack manually, as shown previously. \n",
+    "1. By passing an `nnx.Rngs` object through the `__call__` stack manually, as shown previously.\n",
     "2. By using `nnx.reseed` to set the random state of the model to a specific configuration. This option is less intrusive and can be used even if the model is not designed to enable manual control over the random state.\n",
     "\n",
     "`nnx.reseed` is a function that accepts an arbitrary graph node (this includes [pytrees](https://jax.readthedocs.io/en/latest/working-with-pytrees.html#working-with-pytrees) of `nnx.Module`s) and some keyword arguments containing the new seed or key value for the `nnx.RngStream`s specified by the argument names. `nnx.reseed` will then traverse the graph and update the random state of the matching `nnx.RngStream`s, this includes both setting the `key` to a possibly new value and resetting the `count` to zero.\n",

docs_nnx/guides/randomness.md

@@ -10,26 +10,35 @@ jupytext:
 
 # Randomness
 
-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.Variable`s; and 3) held by the models themselves.
+Flax NNX uses the stateful `nnx.Rngs` class to simplify Jax's handling of random states. For example, the code below uses a `nnx.Rngs` object to define a simple linear model with dropout:
+
+```{code-cell} ipython3
+from flax import nnx
+
+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 always pass `nnx.Rngs` objects to models at initialization (to initialize parameters). For models with nondeterministic outputs like the one above, we also pass `nnx.Rngs` objects to the model's `__call__` method.
 
 The Flax NNX [pseudorandom number generator (PRNG)](https://flax.readthedocs.io/en/latest/glossary.html#term-RNG-sequences) 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](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html), which is `(path + order)`-based, and uses static counters.
-
 > **Note:** To learn more about random number generation in JAX, the `jax.random` API, and PRNG-generated sequences, check out this [JAX PRNG tutorial](https://jax.readthedocs.io/en/latest/random-numbers.html).
 
-Let’ start with some necessary imports:
-
-```{code-cell} ipython3
-from flax import nnx
-import jax
-from jax import random, numpy as jnp
-```
-
 ## `Rngs`, `RngStream`, and `RngState`
 
 In Flax NNX, the `nnx.Rngs` type is the primary convenience API for managing the random state(s). Following Flax Linen's footsteps, `nnx.Rngs` have the ability to create multiple named PRNG key [streams](https://jax.readthedocs.io/en/latest/jep/263-prng.html), each with its own state, for the purpose of having tight control over randomness in the context of [JAX transformations (transforms)](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations).
@@ -85,6 +94,7 @@ dropout = nnx.Dropout(0.5)
 ```
 
 ```{code-cell} ipython3
+import jax.numpy as jnp
 dropout(jnp.ones(4), rngs=rngs)
 ```
 
@@ -127,6 +137,7 @@ As shown above, a PRNG key from the `default` stream can also be generated by ca
 Since a very common pattern is to sample a key and immediately pass it to a function from `jax.random`, both `Rngs` and `RngStream` expose the same functions as methods with the same signature except they don't require a key:
 
 ```{code-cell} ipython3
+import jax
 rngs = nnx.Rngs(0, params=1)
 
 # using jax.random
@@ -224,7 +235,7 @@ In Haiku and Flax Linen, random states are explicitly passed to `Module.apply` e
 
 In Flax NNX, there are two ways to approach this:
 
-1. By passing an `nnx.Rngs` object through the `__call__` stack manually, as shown previously. 
+1. By passing an `nnx.Rngs` object through the `__call__` stack manually, as shown previously.
 2. By using `nnx.reseed` to set the random state of the model to a specific configuration. This option is less intrusive and can be used even if the model is not designed to enable manual control over the random state.
 
 `nnx.reseed` is a function that accepts an arbitrary graph node (this includes [pytrees](https://jax.readthedocs.io/en/latest/working-with-pytrees.html#working-with-pytrees) of `nnx.Module`s) and some keyword arguments containing the new seed or key value for the `nnx.RngStream`s specified by the argument names. `nnx.reseed` will then traverse the graph and update the random state of the matching `nnx.RngStream`s, this includes both setting the `key` to a possibly new value and resetting the `count` to zero.