Vincent Roulet

Hello @gjhuizing, Thanks for catching this! If you are willing to do such a PR that would be great!

Hello @froody, Good catch. The behavior of `adaptive_grad_clip` hides indeed some logic that could mislead users indeed. If you are willing to do a pr to let this function handle...

`clip_by_global_norm` clips but do not necessarily normalize (if the updates are less than clip norm, then they are just returned as is). In other words clip projects on a ball...

Thank you @amosyou for looking into this! I've just tried removing :members: from the LookAheadState and the `fast_state` and `steps_since_sync` were still in the doc. One can even still have...

Ok, but I may have missed something, @amosyou can you try again on your side and see if that works? Your idea of removing the :members: sounded great :). I...

Strange. All states are missing in that screenshot (they're in pink not in blue). There may be another issue? Here is the modification of the code on my side: ```...

Yes, typically, MultiStepsState should be reformatted with attributes instead of fields. Thanks again for looking into that!

> also for ApplyIfFiniteState, I noticed that all the fields are typed as Any. Did we figure out a solution for the typing jnp.array? They should be `Union[jax.Array, int]` for...

> my bad! commenting in rst is not with # that was why they were missing 🤥 > > on a side note, I think to have proper rendering of...

Thanks for looking into this @carlosgmartin. @q-berthet is contributing to a similar approach in #827 (which should be merged soon). Also the whole stochastic gradient estimator part of the codebase...