Merge pull request #4157 from google:nnx-improve-scan

PiperOrigin-RevId: 670690186

docs/nnx/haiku_linen_vs_nnx.rst

@@ -1209,7 +1209,7 @@ in ``__init__`` to scan over the sequence.
       scan_fn = lambda carry, cell, x: cell(carry, x)
       carry = self.cell.initial_state(x.shape[0])
       carry, y = nnx.scan(
-        scan_fn, in_axes=(nnx.Carry, None, 1), out_axes=1
+        scan_fn, in_axes=(nnx.Carry, None, 1), out_axes=(nnx.Carry, 1)
       )(carry, self.cell, x)
 
       return y

docs/nnx/nnx_basics.ipynb

@@ -375,7 +375,7 @@
     "1. The `create_model` function takes in a key and returns an `MLP` object, since we create 5 keys\n",
     "  and use `nnx.vmap` over `create_model` a stack of 5 `MLP` objects is created.\n",
     "2. We use `nnx.scan` to iteratively apply each `MLP` in the stack to the input `x`.\n",
-    "3. The `nnx.scan` API (consciously) deviates from `jax.lax.scan` and instead mimicks `vmap` which is\n",
+    "3. The `nnx.scan` API (consciously) deviates from `jax.lax.scan` and instead mimics `vmap` which is\n",
     "  more expressive. `nnx.scan` allows specifying multiple inputs, the scan axes of each input/output,\n",
     "  and the position of the carry.\n",
     "4. State updates for the `BatchNorm` and `Dropout` layers are automatically propagated\n",
@@ -415,13 +415,13 @@
     "keys = jax.random.split(jax.random.key(0), 5)\n",
     "model = create_model(keys)\n",
     "\n",
-    "@nnx.scan(in_axes=(nnx.Carry, 0), out_axes=0)\n",
-    "def forward(x, model: MLP):\n",
+    "@nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)\n",
+    "def forward(model: MLP, x):\n",
     "  x = model(x)\n",
-    "  return x, None\n",
+    "  return x\n",
     "\n",
     "x = jnp.ones((3, 10))\n",
-    "y, _ = forward(x, model)\n",
+    "y = forward(model, x)\n",
     "\n",
     "print(f'{y.shape = }')\n",
     "nnx.display(model)"

docs/nnx/nnx_basics.md

@@ -229,7 +229,7 @@ Notice the following:
 1. The `create_model` function takes in a key and returns an `MLP` object, since we create 5 keys
   and use `nnx.vmap` over `create_model` a stack of 5 `MLP` objects is created.
 2. We use `nnx.scan` to iteratively apply each `MLP` in the stack to the input `x`.
-3. The `nnx.scan` API (consciously) deviates from `jax.lax.scan` and instead mimicks `vmap` which is
+3. The `nnx.scan` API (consciously) deviates from `jax.lax.scan` and instead mimics `vmap` which is
   more expressive. `nnx.scan` allows specifying multiple inputs, the scan axes of each input/output,
   and the position of the carry.
 4. State updates for the `BatchNorm` and `Dropout` layers are automatically propagated
@@ -243,13 +243,13 @@ def create_model(key: jax.Array):
 keys = jax.random.split(jax.random.key(0), 5)
 model = create_model(keys)
 
-@nnx.scan(in_axes=(nnx.Carry, 0), out_axes=0)
-def forward(x, model: MLP):
+@nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)
+def forward(model: MLP, x):
   x = model(x)
-  return x, None
+  return x
 
 x = jnp.ones((3, 10))
-y, _ = forward(x, model)
+y = forward(model, x)
 
 print(f'{y.shape = }')
 nnx.display(model)