google-deepmind / optax

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

Added mathematical documentation to AdaMax #918

Closed hmludwig closed 5 months ago

hmludwig commented 5 months ago

I added mathematical documentation to AdaMax. Results look like this:

Screenshot 2024-04-09 at 13 17 00
hmludwig commented 5 months ago

On the screenshot v_t appears twice, this isn't the case in the final commit. Just uploaded the wrong screenshot.