Documentation Inconsistency regarding flax.nnx.optimizer.Optimizer

[Lucas-Fernandes-Martins]

Hi everyone,

I believe I found an inconsistency between the implementation of the constructor offlax.nnx.optimizer.Optimizer and what the documentation states.

According to the docs, wrt is supposed to be an optional argument.

However, looking at the class constructor:

@_check_wrt_arg_passed
  def __init__(
    self,
    model: M,
    tx: optax.GradientTransformation,
    *,
    wrt: filterlib.Filter,  # type: ignore
  ):

wrt is actually a required argument.

Ommiting wrt when creating an instance of Optmizer will yield the following error:

TypeError: Missing required argument `wrt`. As of Flax 0.11.0 the `wrt` argument is required, if you want to keep the previous use                    
nnx.ModelAndOptimizer instead of nnx.Optimizer.  

Please let me know if I'm getting something wrong here. Otherwise, I'd be happy to help updating the docs accordingly.

[mohsinm-dev]

Yes, I think you are right. I just checked it. The docs needs to be updated to reflect that wrt is a required parameter, not optional. The README example should also be fixed to include wrt=nnx.Param. The @_check_wrt_arg_passed decorator actually requires wrt and throws the exact error you mentioned

Suggested fixes:

• Update Optimizer.init docstring to mark wrt as required and note it should match the filter used in nnx.grad.

[Lucas-Fernandes-Martins]

Thank you @mohsinm-dev for confirming! I opened this PR to correct the docstring. On the other hand, I wasn't able to find the example in the README which uses the optimizer (I only see the MLP, CNN, and Autoencoder examples). Apologies if I'm looking in the wrong place.

[mohsinm-dev]

Thanks for the PR. I think this change is sufficient to resolve the issue @samanklesaria can you verify this change?

I also noticed a couple of spots that still show the old usage:

• flax/nnx/README.md: shows pre-0.11 style without wrt and with optimizer.update(grads). It should be:
  • optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
  • optimizer.update(model, grads)
• docs_nnx/mnist_tutorial.md: already uses wrt=nnx.Param but still has optimizer.update(grads); should be
  optimizer.update(model, grads).
• For reference, the migration guide is already correct: docs_nnx/migrating/nnx_010_to_nnx_011.rst (uses wrt=nnx.Param and update(model, grads)).

[samanklesaria]

You're right. The README and docs should be updated to show that wrt is required. @mohsinm-dev : the change to docs_nnx/mnist_tutorial.md you recommend is already made in #5002- it should be merged soon.