feat: add `scale_by_gradient_norm`

[SauravMaheshkar]

Reference: #594

This PR aims to add a new GradientTransformation which allows to scale by the gradient norm.

Request for Review: @fabianp

[SauravMaheshkar]

CI is green again 😄

[SauravMaheshkar]

Gentle ping: @fabianp

[fabianp]

thanks for the ping, will look into this ASAP

[fabianp]

Hi @SauravMaheshkar ,

a couple more comments after discussing this issue with the dev team:

• let's promote this to the main package, since normalized GD is a fairly classical method :-)
• we found the name scale_by_gradient_norm a bit misleading, as it's really scaled by the inverse of the gradient norm. A better name for the method would (IMO) be normalize_by_update_norm.

what do you think about these comments? do they sound reasonable?

[SauravMaheshkar]

Hi @SauravMaheshkar ,

a couple more comments after discussing this issue with the dev team:

• let's promote this to the main package, since normalized GD is a fairly classical method :-)
• we found the name scale_by_gradient_norm a bit misleading, as it's really scaled by the inverse of the gradient norm. A better name for the method would (IMO) be normalize_by_update_norm.

what do you think about these comments? do they sound reasonable?

Hey @fabianp glad to know there's interest. That sounds reasonable.

Let me do some refactoring. Since it's going into main, I'll put some more effort in the docstrings as well. Will push soon 😄

[fabianp]

amazing, thanks!

[SauravMaheshkar]

@fabianp I've refactored my branch to add the transform to main rather than contrib. Request for Review 😄

[SauravMaheshkar]

I've added a doctest 👍🏽 @fabianp

[fabianp]

looks good to me, but I forgot to ask you one (hopefully) last thing:

please add this new function to the API documentation: https://github.com/google-deepmind/optax/blob/main/docs/api/transformations.rst

thanks! 🙏🏼

[SauravMaheshkar]

looks good to me, but I forgot to ask you one (hopefully) last thing:

please add this new function to the API documentation: https://github.com/google-deepmind/optax/blob/main/docs/api/transformations.rst

thanks! 🙏🏼

Added in 5c01f6f

[fabianp]

thanks @SauravMaheshkar !

[fabianp]

duplicated https://github.com/google-deepmind/optax/pull/1000 since I'm having some issues getting this merged due to internal errors

[SauravMaheshkar]

duplicated #1000 since I'm having some issues getting this merged due to internal errors

Typically copybara I assume 😅

[fabianp]

yup. Finally merged 🎉

[SauravMaheshkar]

yup. Finally merged 🎉

Thanks a ton. Feels good to have it merged. Now on to #652