Merge pull request #5061 from Lucas-Fernandes-Martins:main

PiperOrigin-RevId: 826201674

Author: Flax Authors

flax/nnx/README.md

@@ -39,7 +39,7 @@ class Model(nnx.Module):
 
 
 model = Model(2, 64, 3, rngs=nnx.Rngs(0))  # eager initialization
-optimizer = nnx.Optimizer(model, optax.adam(1e-3))  # reference sharing
+optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)  # reference sharing
 
 @nnx.jit  # automatic state management
 def train_step(model, optimizer, x, y):
@@ -48,7 +48,7 @@ def train_step(model, optimizer, x, y):
     return ((y_pred - y) ** 2).mean()
 
   loss, grads = nnx.value_and_grad(loss_fn)(model)
-  optimizer.update(grads)  # inplace updates
+  optimizer.update(model, grads)  # inplace updates
 
   return loss
 ```

flax/nnx/training/optimizer.py

@@ -144,12 +144,12 @@ def __init__(
     Args:
       model: An NNX Module.
       tx: An Optax gradient transformation.
-      wrt: optional argument to filter for which :class:`Variable`'s to keep
+      wrt: filter to specify for which :class:`Variable`'s to keep
         track of in the optimizer state. These should be the :class:`Variable`'s
         that you plan on updating; i.e. this argument value should match the
         ``wrt``  argument passed to the ``nnx.grad`` call that will generate the
         gradients that will be passed into the ``grads`` argument of the
-        :func:`update` method.
+        :func:`update` method. The filter should match the filter used in nnx.grad.
     """
     if isinstance(wrt, _Missing):
       raise TypeError(