diff --git a/docs_nnx/guides/randomness.ipynb b/docs_nnx/guides/randomness.ipynb
index ad7c16c92..2d8cd33bf 100644
--- a/docs_nnx/guides/randomness.ipynb
+++ b/docs_nnx/guides/randomness.ipynb
@@ -11,11 +11,13 @@
},
{
"cell_type": "code",
- "execution_count": 2,
+ "execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
+ "import jax\n",
"from flax import nnx\n",
+ "import jax.numpy as jnp\n",
"\n",
"class Model(nnx.Module):\n",
" def __init__(self, *, rngs: nnx.Rngs):\n",
@@ -64,13 +66,13 @@
},
{
"cell_type": "code",
- "execution_count": 3,
+ "execution_count": 2,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -82,7 +84,7 @@
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -93,7 +95,7 @@
}
],
"source": [
- "rngs = nnx.Rngs(params=0, dropout=random.key(1))\n",
+ "rngs = nnx.Rngs(params=0, dropout=jax.random.key(1))\n",
"nnx.display(rngs)"
]
},
@@ -108,13 +110,13 @@
},
{
"cell_type": "code",
- "execution_count": 4,
+ "execution_count": 26,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -126,7 +128,7 @@
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -161,7 +163,7 @@
},
{
"cell_type": "code",
- "execution_count": 5,
+ "execution_count": 27,
"metadata": {},
"outputs": [],
"source": [
@@ -184,16 +186,16 @@
},
{
"cell_type": "code",
- "execution_count": 6,
+ "execution_count": 28,
"metadata": {},
"outputs": [],
"source": [
- "dropout = nnx.Dropout(0.5, deterministic=False)"
+ "dropout = nnx.Dropout(0.5)"
]
},
{
"cell_type": "code",
- "execution_count": 7,
+ "execution_count": 29,
"metadata": {},
"outputs": [
{
@@ -202,13 +204,12 @@
"Array([2., 0., 2., 2.], dtype=float32)"
]
},
- "execution_count": 7,
+ "execution_count": 29,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
- "import jax.numpy as jnp\n",
"dropout(jnp.ones(4), rngs=rngs)"
]
},
@@ -243,13 +244,13 @@
},
{
"cell_type": "code",
- "execution_count": 8,
+ "execution_count": 30,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -261,7 +262,7 @@
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -301,11 +302,10 @@
},
{
"cell_type": "code",
- "execution_count": 9,
+ "execution_count": 31,
"metadata": {},
"outputs": [],
"source": [
- "import jax\n",
"rngs = nnx.Rngs(0, params=1)\n",
"\n",
"# using jax.random\n",
@@ -321,26 +321,26 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Forking random state"
+ "## Split random state"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "Say you want to train a model that uses dropout on a batch of data. You don't want to use the same random state for every dropout mask in your batch. Instead, you want to fork the random state into separate pieces for each layer. This can be accomplished with the `fork` method, as shown below."
+ "Say you want to train a model that uses dropout on a batch of data. You don't want to use the same random state for every dropout mask in your batch. Instead, you want to split the random state into separate pieces for each layer. This can be accomplished with the `split` method, as shown below."
]
},
{
"cell_type": "code",
- "execution_count": 10,
+ "execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"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, deterministic=False)\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))"
@@ -348,7 +348,7 @@
},
{
"cell_type": "code",
- "execution_count": 11,
+ "execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
@@ -357,7 +357,7 @@
},
{
"cell_type": "code",
- "execution_count": 12,
+ "execution_count": 34,
"metadata": {},
"outputs": [],
"source": [
@@ -368,7 +368,7 @@
},
{
"cell_type": "code",
- "execution_count": 13,
+ "execution_count": 35,
"metadata": {},
"outputs": [
{
@@ -403,20 +403,20 @@
" ))"
]
},
- "execution_count": 13,
+ "execution_count": 35,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"dropout_rngs = nnx.Rngs(1)\n",
- "forked_rngs = dropout_rngs.fork(split=5)\n",
- "(dropout_rngs, forked_rngs)"
+ "split_rngs = dropout_rngs.split(5)\n",
+ "(dropout_rngs, split_rngs)"
]
},
{
"cell_type": "code",
- "execution_count": 14,
+ "execution_count": 36,
"metadata": {},
"outputs": [
{
@@ -425,20 +425,20 @@
"(5, 10)"
]
},
- "execution_count": 14,
+ "execution_count": 36,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
- "model_forward(model, jnp.ones((5, 20)), forked_rngs).shape"
+ "model_forward(model, jnp.ones((5, 20)), split_rngs).shape"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "The output of `rng.fork` is another `Rng` with keys and counts that have an expanded shape. In the example above, the `RngKey` and `RngCount` of `dropout_rngs` have shape `()`, but in `forked_rngs` they have shape `(5,)`."
+ "The output of `rng.split` is another `Rng` with keys and counts that have an expanded shape. In the example above, the `RngKey` and `RngCount` of `dropout_rngs` have shape `()`, but in `split_rngs` they have shape `(5,)`."
]
},
{
@@ -457,7 +457,7 @@
},
{
"cell_type": "code",
- "execution_count": 18,
+ "execution_count": 5,
"metadata": {},
"outputs": [
{
@@ -472,7 +472,7 @@
"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, deterministic=False, rngs=rngs)\n",
+ " self.drop = nnx.Dropout(0.1, rngs=rngs)\n",
"\n",
" def __call__(self, x):\n",
" return nnx.relu(self.drop(self.linear(x)))\n",
@@ -501,13 +501,13 @@
},
{
"cell_type": "code",
- "execution_count": 21,
+ "execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -519,7 +519,7 @@
{
"data": {
"text/html": [
- "
"
+ "
"
],
"text/plain": [
""
@@ -560,7 +560,7 @@
},
{
"cell_type": "code",
- "execution_count": 22,
+ "execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
@@ -581,14 +581,28 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Forking implicit random state\n",
+ "## Splitting implicit random state\n",
+ "\n",
+ "We saw above how to use `rng.split` when passing explicit random state through [Flax NNX transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) like `nnx.vmap` or `nnx.pmap`. We can do the same for implicit random state using the [View](https://flax.readthedocs.io/en/latest/guides/view.html) `nnx.with_rngs`. Consider the example below, which generates a batch of samples from the nondeterministic model we defined above."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "in_axes = nnx.prefix(model, {'dropout': 0, ...: None}, graph=False)\n",
"\n",
- "We saw above how to use `rng.fork` when passing explicit random state through [Flax NNX transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) like `nnx.vmap` or `nnx.pmap`. The decorator `nnx.fork_rngs` allows this for implicit random state. Consider the example below, which generates a batch of samples from the nondeterministic model we defined above."
+ "@nnx.with_rngs(split=5, only='dropout', graph=False)\n",
+ "@nnx.vmap(in_axes=(in_axes, None), out_axes=0, graph=False)\n",
+ "def sample_from_model(model, x):\n",
+ " return model(x)"
]
},
{
"cell_type": "code",
- "execution_count": 23,
+ "execution_count": 11,
"metadata": {},
"outputs": [
{
@@ -600,13 +614,6 @@
}
],
"source": [
- "rng_axes = nnx.StateAxes({'dropout': 0, ...: None})\n",
- "\n",
- "@nnx.fork_rngs(split={'dropout': 5})\n",
- "@nnx.vmap(in_axes=(rng_axes, None), out_axes=0)\n",
- "def sample_from_model(model, x):\n",
- " return model(x)\n",
- "\n",
"print(sample_from_model(model, x).shape)"
]
},
@@ -614,9 +621,8 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "Here `sample_from_model` is modified by two decorators:\n",
- "- The function we get from the `nnx.vmap` decorator expects that the random state of the `model` argument has already been split into 5 pieces. It runs the model once for each random key.\n",
- "- The function we get from the `nnx.fork_rngs` decorator splits the random state of its `model` argument into five pieces before passing it on to the inner function."
+ "The call to the decorator `with_rngs` with the `split` argument makes `sample_from_model` split the rngs of the `model` argument before the function body is applied. \n",
+ "Within the function body, we tell `vmap` to only map over values within Variables tagged as \"dropout\" using `nnx.prefix` to specify the `in_axes`."
]
},
{
@@ -642,7 +648,7 @@
},
{
"cell_type": "code",
- "execution_count": 24,
+ "execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
@@ -651,14 +657,14 @@
"class RNNCell(nnx.Module):\n",
" def __init__(self, din, dout, rngs):\n",
" self.linear = nnx.Linear(dout + din, dout, rngs=rngs)\n",
- " self.drop = nnx.Dropout(0.1, deterministic=False, rngs=rngs, rng_collection='recurrent_dropout')\n",
+ " self.drop = nnx.Dropout(0.1, rngs=rngs, rng_collection='recurrent_dropout')\n",
" self.dout = dout\n",
" self.count = Count(jnp.array(0, jnp.uint32))\n",
"\n",
" def __call__(self, h, x) -> tuple[jax.Array, jax.Array]:\n",
" h = self.drop(h) # Recurrent dropout.\n",
" y = nnx.relu(self.linear(jnp.concatenate([h, x], axis=-1)))\n",
- " self.count.value += 1\n",
+ " self.count[...] = self.count[...] + 1\n",
" return y, y\n",
"\n",
" def initial_state(self, batch_size: int):\n",
@@ -672,13 +678,13 @@
"metadata": {},
"source": [
"Next, use `nnx.scan` over an `unroll` function to implement the `rnn_forward` operation:\n",
- "- The key ingredient of recurrent dropout is to apply the same dropout mask across all time steps. Therefore, to achieve this you will pass `nnx.StateAxes` to `nnx.scan`'s `in_axes`, specifying that the `cell`'s `recurrent_dropout` PRNG stream will be broadcast, and the rest of the `RNNCell`'s state will be carried over.\n",
- "- Also, the hidden state `h` will be the `nnx.scan`'s `Carry` variable, and the sequence `x` will be `scan`ned over its axis `1`."
+ "- The key ingredient of recurrent dropout is to apply the same dropout mask across all time steps. To achieve this, you must use `nnx.with_rngs` with the `broadcast` argument. This will broadcast the `recurrent_dropout` stream within the implicit state. To specify that this state should traversed in the call to `nnx.scan`, you can use a call to `nnx.prefix` to generate the appropriate `in_axes` argument. \n",
+ "- The hidden state `h` will be the `nnx.scan`'s `Carry` variable, and the sequence `x` will be `scan`ned over its axis `1`."
]
},
{
"cell_type": "code",
- "execution_count": 25,
+ "execution_count": 16,
"metadata": {},
"outputs": [
{
@@ -686,36 +692,36 @@
"output_type": "stream",
"text": [
"y.shape = (4, 20, 16)\n",
- "cell.count.value = Array(20, dtype=uint32)\n"
+ "cell.count[...] = Array(20, dtype=uint32)\n"
]
}
],
"source": [
- "@nnx.jit\n",
+ "@nnx.jit(graph=False)\n",
"def rnn_forward(cell: RNNCell, x: jax.Array):\n",
- " h = cell.initial_state(batch_size=x.shape[0])\n",
- "\n",
- " # Broadcast the 'recurrent_dropout' PRNG state to have the same mask on every step.\n",
- " state_axes = nnx.StateAxes({'recurrent_dropout': None, ...: nnx.Carry})\n",
- " @nnx.scan(in_axes=(state_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1))\n",
- " def unroll(cell: RNNCell, h, x) -> tuple[jax.Array, jax.Array]:\n",
- " h, y = cell(h, x)\n",
- " return h, y\n",
- "\n",
- " h, y = unroll(cell, h, x)\n",
- " return y\n",
+ " h = cell.initial_state(batch_size=x.shape[0])\n",
+ " # Broadcast the 'recurrent_dropout' PRNG state to have the same mask on every step.\n",
+ " broadcast_cell = nnx.with_rngs(cell, broadcast=x.shape[1], graph=False)\n",
+ " in_axes = nnx.prefix(broadcast_cell, {'recurrent_dropout': 0, ...: None}, graph=False)\n",
+ " @nnx.scan(in_axes=(in_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1), graph=False)\n",
+ " def unroll(cell, h, x) -> tuple[jax.Array, jax.Array]:\n",
+ " h, y = cell(h, x)\n",
+ " return h, y\n",
+ " _, y = unroll(broadcast_cell, h, x)\n",
+ " return y\n",
"\n",
"x = jnp.ones((4, 20, 8))\n",
"y = rnn_forward(cell, x)\n",
"\n",
"print(f'{y.shape = }')\n",
- "print(f'{cell.count.value = }')"
+ "print(f'{cell.count[...] = }')"
]
}
],
"metadata": {
"jupytext": {
- "formats": "ipynb,md:myst"
+ "formats": "ipynb,md",
+ "main_language": "python"
},
"language_info": {
"codemirror_mode": {
@@ -727,7 +733,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.11.6"
+ "version": "3.13.5"
}
},
"nbformat": 4,
diff --git a/docs_nnx/guides/randomness.md b/docs_nnx/guides/randomness.md
index 3a3881a27..2d3df8087 100644
--- a/docs_nnx/guides/randomness.md
+++ b/docs_nnx/guides/randomness.md
@@ -1,19 +1,23 @@
---
-jupytext:
- formats: ipynb,md:myst
- text_representation:
- extension: .md
- format_name: myst
- format_version: 0.13
- jupytext_version: 1.13.8
+jupyter:
+ jupytext:
+ formats: ipynb,md
+ main_language: python
+ text_representation:
+ extension: .md
+ format_name: markdown
+ format_version: '1.3'
+ jupytext_version: 1.13.8
---
# Randomness
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
+```python
+import jax
from flax import nnx
+import jax.numpy as jnp
class Model(nnx.Module):
def __init__(self, *, rngs: nnx.Rngs):
@@ -55,8 +59,8 @@ To create an `nnx.Rngs` object you can simply pass an integer seed or `jax.rando
Here's an example:
-```{code-cell} ipython3
-rngs = nnx.Rngs(params=0, dropout=random.key(1))
+```python
+rngs = nnx.Rngs(params=0, dropout=jax.random.key(1))
nnx.display(rngs)
```
@@ -64,7 +68,7 @@ Notice that the `key` and `count` `nnx.Variable`s contain the PRNG key stream na
To generate new keys, you can access one of the streams and use its `__call__` method with no arguments. This will return a new key by using `random.fold_in` with the current `key` and `count`. The `count` is then incremented so that subsequent calls will return new keys.
-```{code-cell} ipython3
+```python
params_key = rngs.params()
dropout_key = rngs.dropout()
@@ -73,34 +77,30 @@ nnx.display(rngs)
Note that the `key` attribute does not change when new PRNG keys are generated.
-+++
### Using random state with flax Modules.
Almost all flax Modules require a random state for initialization. In a `Linear` layer, for example, we need to sample the weights and biases from the appropriate Normal distribution. Random state is provided using the `rngs` keyword argument at initialization.
-```{code-cell} ipython3
+```python
linear = nnx.Linear(20, 10, rngs=rngs)
```
Specifically, this will use the RngSteam `rngs.params` for weight initialization. The `params` stream is also used for initialization of `nnx.Conv`, `nnx.ConvTranspose`, and `nnx.Embed`.
-+++
The `nnx.Dropout` module also requires a random state, but it requires this state at *call* time rather than initialization. Once again, we can pass it random state using the `rngs` keyword argument.
-```{code-cell} ipython3
-dropout = nnx.Dropout(0.5, deterministic=False)
+```python
+dropout = nnx.Dropout(0.5)
```
-```{code-cell} ipython3
-import jax.numpy as jnp
+```python
dropout(jnp.ones(4), rngs=rngs)
```
The `nnx.Dropout` layer will use the rng's `dropout` stream. This also applies to Modules that use `Dropout` as a sub-Module, like `nnx.MultiHeadAttention`.
-+++
To summarize, there are only two standard PRNG key stream names used by Flax NNX's built-in layers, shown in the table below:
@@ -109,14 +109,13 @@ To summarize, there are only two standard PRNG key stream names used by Flax NNX
| `params` | Used for parameter initialization |
| `dropout` | Used by `nnx.Dropout` to create dropout masks |
-+++
### Default PRNG key stream
One of the downsides of having named streams is that the user needs to know all the possible names that a model will use when creating the `nnx.Rngs` object. While this could be solved with some documentation, Flax NNX provides a `default` stream that can be
be used as a fallback when a stream is not found. To use the default PRNG key stream, you can simply pass an integer seed or `jax.random.key` as the first positional argument.
-```{code-cell} ipython3
+```python
rngs = nnx.Rngs(0, params=1)
key1 = rngs.params() # Call params.
@@ -131,13 +130,11 @@ As shown above, a PRNG key from the `default` stream can also be generated by ca
> **Note**
>
For large projects it is recommended to use named streams to avoid potential conflicts. For small projects or quick prototyping just using the `default` stream is a good choice.
-+++
### jax.random shorthand methods
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
+```python
rngs = nnx.Rngs(0, params=1)
# using jax.random
@@ -149,57 +146,54 @@ z1 = rngs.normal((2, 3)) # generates key from rngs.default
z2 = rngs.params.bernoulli(0.5, (10,)) # generates key from rngs.params
```
-## Forking random state
+## Split random state
-+++
-Say you want to train a model that uses dropout on a batch of data. You don't want to use the same random state for every dropout mask in your batch. Instead, you want to fork the random state into separate pieces for each layer. This can be accomplished with the `fork` method, as shown below.
+Say you want to train a model that uses dropout on a batch of data. You don't want to use the same random state for every dropout mask in your batch. Instead, you want to split the random state into separate pieces for each layer. This can be accomplished with the `split` method, as shown below.
-```{code-cell} ipython3
+```python
class Model(nnx.Module):
def __init__(self, rngs: nnx.Rngs):
self.linear = nnx.Linear(20, 10, rngs=rngs)
- self.drop = nnx.Dropout(0.1, deterministic=False)
+ self.drop = nnx.Dropout(0.1)
def __call__(self, x, rngs):
return nnx.relu(self.drop(self.linear(x), rngs=rngs))
```
-```{code-cell} ipython3
+```python
model = Model(rngs=nnx.Rngs(0))
```
-```{code-cell} ipython3
+```python
@nnx.vmap(in_axes=(None, 0, 0), out_axes=0)
def model_forward(model, x, rngs):
return model(x, rngs=rngs)
```
-```{code-cell} ipython3
+```python
dropout_rngs = nnx.Rngs(1)
-forked_rngs = dropout_rngs.fork(split=5)
-(dropout_rngs, forked_rngs)
+split_rngs = dropout_rngs.split(5)
+(dropout_rngs, split_rngs)
```
-```{code-cell} ipython3
-model_forward(model, jnp.ones((5, 20)), forked_rngs).shape
+```python
+model_forward(model, jnp.ones((5, 20)), split_rngs).shape
```
-The output of `rng.fork` is another `Rng` with keys and counts that have an expanded shape. In the example above, the `RngKey` and `RngCount` of `dropout_rngs` have shape `()`, but in `forked_rngs` they have shape `(5,)`.
+The output of `rng.split` is another `Rng` with keys and counts that have an expanded shape. In the example above, the `RngKey` and `RngCount` of `dropout_rngs` have shape `()`, but in `split_rngs` they have shape `(5,)`.
-+++
# Implicit Random State
-+++
So far, we have looked at passing random state directly to each Module when it gets called. But there's another way to handle call-time randomness in flax: we can bundle the random state into the Module itself. This makes the random state is just another type of Module state. Using implicit random state requires passing the `rngs` keyward argument when initializing the module rather than when calling it. For example, here is how we might construct the simple `Module` we defined earlier using an implicit style.
-```{code-cell} ipython3
+```python
class Model(nnx.Module):
def __init__(self, rngs: nnx.Rngs):
self.linear = nnx.Linear(20, 10, rngs=rngs)
- self.drop = nnx.Dropout(0.1, deterministic=False, rngs=rngs)
+ self.drop = nnx.Dropout(0.1, rngs=rngs)
def __call__(self, x):
return nnx.relu(self.drop(self.linear(x)))
@@ -212,13 +206,12 @@ print(f'{y.shape = }')
This implicit state handling style is less verbose than passing RNGs explicitly, and more closely resembles code in other deep learning frameworks like PyTorch. However, as we'll see in the following sections, using implicit state makes it less obvious how to apply jax transformations to your Modules. With explicit state, you can usually use tranforms like `jax.vmap` directly. With implicit state, you'll need to some extra tricks with `nnx.vmap` to make everything work. Because of this additional complexity, we recommend that new flax projects stick to the explicit style.
-+++
## Filtering random state
Implicit random state can be manipulated using [Filters](https://flax.readthedocs.io/en/latest/guides/filters_guide.html) just like any other type of state. It can be filtered using types (`nnx.RngState`, `nnx.RngKey`, `nnx.RngCount`) or using strings corresponding to the stream names (refer to [the Flax NNX `Filter` DSL](https://flax.readthedocs.io/en/latest/guides/filters_guide.html#the-filter-dsl)). Here's an example using `nnx.state` with various filters to select different substates of the `Rngs` inside a `Model`:
-```{code-cell} ipython3
+```python
model = Model(nnx.Rngs(params=0, dropout=1))
rng_state = nnx.state(model, nnx.RngState) # All random states.
@@ -242,7 +235,7 @@ In Flax NNX, there are two ways to approach this:
Here's an example of how to use `nnx.reseed` to reset the random state of the `nnx.Dropout` layer and verify that the computation is identical to the first time the model was called:
-```{code-cell} ipython3
+```python
model = Model(nnx.Rngs(params=0, dropout=1))
x = jnp.ones((1, 20))
@@ -256,30 +249,29 @@ assert not jnp.allclose(y1, y2) # different
assert jnp.allclose(y1, y3) # same
```
-## Forking implicit random state
+## Splitting implicit random state
-We saw above how to use `rng.fork` when passing explicit random state through [Flax NNX transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) like `nnx.vmap` or `nnx.pmap`. The decorator `nnx.fork_rngs` allows this for implicit random state. Consider the example below, which generates a batch of samples from the nondeterministic model we defined above.
+We saw above how to use `rng.split` when passing explicit random state through [Flax NNX transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) like `nnx.vmap` or `nnx.pmap`. We can do the same for implicit random state using the [View](https://flax.readthedocs.io/en/latest/guides/view.html) `nnx.with_rngs`. Consider the example below, which generates a batch of samples from the nondeterministic model we defined above.
-```{code-cell} ipython3
-rng_axes = nnx.StateAxes({'dropout': 0, ...: None})
+```python
+in_axes = nnx.prefix(model, {'dropout': 0, ...: None}, graph=False)
-@nnx.fork_rngs(split={'dropout': 5})
-@nnx.vmap(in_axes=(rng_axes, None), out_axes=0)
+@nnx.with_rngs(split=5, only='dropout', graph=False)
+@nnx.vmap(in_axes=(in_axes, None), out_axes=0, graph=False)
def sample_from_model(model, x):
return model(x)
+```
+```python
print(sample_from_model(model, x).shape)
```
-Here `sample_from_model` is modified by two decorators:
-- The function we get from the `nnx.vmap` decorator expects that the random state of the `model` argument has already been split into 5 pieces. It runs the model once for each random key.
-- The function we get from the `nnx.fork_rngs` decorator splits the random state of its `model` argument into five pieces before passing it on to the inner function.
+The call to the decorator `with_rngs` with the `split` argument makes `sample_from_model` split the rngs of the `model` argument before the function body is applied.
+Within the function body, we tell `vmap` to only map over values within Variables tagged as "dropout" using `nnx.prefix` to specify the `in_axes`.
-+++
## Transforming implicit state
-+++
In the previous section, we showed how to use `nnx.vmap` with a module that contained implicit random state. But we can use other `nnx` transformations too! Remember: implicit random state isn't different from any other type of Model state, and this applies to Flax NNX transforms too. This means you can use the Flax NNX state handling APIs of each transform to get the results you want. For a more involved example, let’s explore how to implement recurrent dropout on an `RNNCell` using `nnx.scan`.
@@ -290,20 +282,20 @@ We'll start by constructing the `RNNCell` class:
- Then, define an `initial_state` function to create the initial state of the `RNNCell`.
- Finally, instantiate `RNNCell`.
-```{code-cell} ipython3
+```python
class Count(nnx.Variable): pass
class RNNCell(nnx.Module):
def __init__(self, din, dout, rngs):
self.linear = nnx.Linear(dout + din, dout, rngs=rngs)
- self.drop = nnx.Dropout(0.1, deterministic=False, rngs=rngs, rng_collection='recurrent_dropout')
+ self.drop = nnx.Dropout(0.1, rngs=rngs, rng_collection='recurrent_dropout')
self.dout = dout
self.count = Count(jnp.array(0, jnp.uint32))
def __call__(self, h, x) -> tuple[jax.Array, jax.Array]:
h = self.drop(h) # Recurrent dropout.
y = nnx.relu(self.linear(jnp.concatenate([h, x], axis=-1)))
- self.count.value += 1
+ self.count[...] = self.count[...] + 1
return y, y
def initial_state(self, batch_size: int):
@@ -313,27 +305,26 @@ cell = RNNCell(8, 16, nnx.Rngs(params=0, recurrent_dropout=1))
```
Next, use `nnx.scan` over an `unroll` function to implement the `rnn_forward` operation:
-- The key ingredient of recurrent dropout is to apply the same dropout mask across all time steps. Therefore, to achieve this you will pass `nnx.StateAxes` to `nnx.scan`'s `in_axes`, specifying that the `cell`'s `recurrent_dropout` PRNG stream will be broadcast, and the rest of the `RNNCell`'s state will be carried over.
-- Also, the hidden state `h` will be the `nnx.scan`'s `Carry` variable, and the sequence `x` will be `scan`ned over its axis `1`.
+- The key ingredient of recurrent dropout is to apply the same dropout mask across all time steps. To achieve this, you must use `nnx.with_rngs` with the `broadcast` argument. This will broadcast the `recurrent_dropout` stream within the implicit state. To specify that this state should traversed in the call to `nnx.scan`, you can use a call to `nnx.prefix` to generate the appropriate `in_axes` argument.
+- The hidden state `h` will be the `nnx.scan`'s `Carry` variable, and the sequence `x` will be `scan`ned over its axis `1`.
-```{code-cell} ipython3
-@nnx.jit
+```python
+@nnx.jit(graph=False)
def rnn_forward(cell: RNNCell, x: jax.Array):
- h = cell.initial_state(batch_size=x.shape[0])
-
- # Broadcast the 'recurrent_dropout' PRNG state to have the same mask on every step.
- state_axes = nnx.StateAxes({'recurrent_dropout': None, ...: nnx.Carry})
- @nnx.scan(in_axes=(state_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1))
- def unroll(cell: RNNCell, h, x) -> tuple[jax.Array, jax.Array]:
- h, y = cell(h, x)
- return h, y
-
- h, y = unroll(cell, h, x)
- return y
+ h = cell.initial_state(batch_size=x.shape[0])
+ # Broadcast the 'recurrent_dropout' PRNG state to have the same mask on every step.
+ broadcast_cell = nnx.with_rngs(cell, broadcast=x.shape[1], graph=False)
+ in_axes = nnx.prefix(broadcast_cell, {'recurrent_dropout': 0, ...: None}, graph=False)
+ @nnx.scan(in_axes=(in_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1), graph=False)
+ def unroll(cell, h, x) -> tuple[jax.Array, jax.Array]:
+ h, y = cell(h, x)
+ return h, y
+ _, y = unroll(broadcast_cell, h, x)
+ return y
x = jnp.ones((4, 20, 8))
y = rnn_forward(cell, x)
print(f'{y.shape = }')
-print(f'{cell.count.value = }')
+print(f'{cell.count[...] = }')
```
diff --git a/flax/nnx/rnglib.py b/flax/nnx/rnglib.py
index 1b87c0887..196ddeca2 100644
--- a/flax/nnx/rnglib.py
+++ b/flax/nnx/rnglib.py
@@ -460,7 +460,7 @@ def items(self):
def split(self, k: tp.Mapping[filterlib.Filter, int | tuple[int, ...]] | int | tuple[int, ...]):
"""
- Splits the keys of the newly created ``Rngs`` object.
+ Splits the keys of a ``Rngs`` object.
Example::
>>> from flax import nnx