nnx.Optimizer.update should return the updates #5182

Activity

1. vfdev-5 commented on Jan 12, 2026

   

   vfdev-5

   on Jan 12, 2026

   Collaborator

   More actions

   @am001122 what exactly you would like to return from update?
   I think user can easily inspect opt_state and model. If you need to return new_params or new_opt_state or some other intermediate object, I think the most simple is to make a derived class from nnx.Optmizer and override update method and return desired objects. What do you think?
2. am001122 commented on Jan 13, 2026

   

   am001122

   on Jan 13, 2026

   Author

   More actions

   Hi @vfdev-5 , sorry, I should have been more precise: I meant to literally return the variable updates containing the updates to the weights. As you say, its easy to access nnx.Optimizer.opt_state, or the state of the optimized nnx.Module after the update. But to get the change in weights requires some external bookkeeping to save the previous weights, take the new ones, and do an extra subtraction operation, all of which could be skipped if updates was returned.
3. vfdev-5 commented on Jan 14, 2026

   

   vfdev-5

   on Jan 14, 2026

   Collaborator

   More actions

   @am001122 thanks for the context. I was thinking that users can make a derived class of nnx.Optimizer with customized update method:

   import flax.nnx as nnx
   
   class Optimizer(nnx.Optimizer):
     def update(self, model, grads, **kwargs):
       # copied from nnx.Optimizer.update method
       param_arrays = nnx.pure(nnx.state(model, self.wrt))
       grad_arrays = nnx.pure(nnx.state(grads, self.wrt))
       opt_state_arrays = nnx.pure(self.opt_state)
       kwargs_arrays = nnx.pure(kwargs)
   
       updates, new_opt_state = self.tx.update(
         grad_arrays, opt_state_arrays, param_arrays, **kwargs_arrays
       )
       new_params = optax.apply_updates(param_arrays, updates)
   
       nnx.update(model, new_params)
       nnx.update(self.opt_state, nnx.state(new_opt_state))
       self.step[...] += 1
       
       # we now return updates for other usages
       return updates

   I agree that ideally this could be done in few copied lines.
4. am001122 commented on Jan 14, 2026

   

   am001122

   on Jan 14, 2026

   Author

   More actions

   Definitely an option---I didn't take it because it involves almost completely re-implementing nnx.Optimizer in my code, and I didn't want the risk of drifting out of sync with upstream. My sense was that since nnx.Optimizer.update currently doesn't return anything this would be a non-breaking change, and a change that nicely mirrors the underlying optax API that nnx.Optimizer thinly wraps to handle state updates. But certainly there are at least these two acceptable workarounds, so it's not really an issue.
5. qGentry commented on Jan 22, 2026

   

   qGentry

   on Jan 22, 2026

   More actions

   +1, would love to see the nnx.Optimizer returning updates without having to copy-paste source code from flax to implement it
6. copybara-service

   closed this as completedin #5214on Mar 31, 2026