search

google-deepmind / optax

Optax is a gradient processing and optimization library for JAX.
https://optax.readthedocs.io
Apache License 2.0
1.64k stars 179 forks source link

New Example on creating a custom `GradientTransformation` #652

Open SauravMaheshkar opened 9 months ago

SauravMaheshkar commented 9 months ago

Motivation

Recently for a personal project (#641) I wanted to created a custom GradientTransformation object and found out that there wasn't a tutorial or walkthrough on how to create one. While the "Gradient Transformations" section in the README.md section was somewhat helpful, I still had to look around and search for example implementations in other repositories and within the contrib sub-package.

Proposal

Create a new example or a page in the documentation containing instructions on how to create a custom GradientTransformation object.

Happy to contribute and work on this.

fabianp commented 9 months ago

hey @SauravMaheshkar , yes, I believe a short example on this on the example/ directory would be welcomed!. I'll assign you this issue.

Thanks for suggesting this!

SauravMaheshkar commented 4 months ago

Also tagging #958 and #641