Code

Adds AdeMAMix Optimizer to contrib#1104

Code

Open

mathDR wants to merge 33 commits intogoogle-deepmind:maingoogle-deepmind/optax:mainfrom mathDR:mainmathDR/ademamix_optax:mainCopy head branch name to clipboard

+721

Commits

Commits on Oct 13, 2024

• getting ademamix docs and notebook running

  mathDR

  committed
  daed085

• fixed imports

  mathDR

  committed
  c47261c

Commits on Oct 14, 2024

• fixed linting errors

  mathDR

  committed
  bf2d4a8

• ran notebook in order

  mathDR

  committed
  7cb270a

• Merge branch 'google-deepmind:main' into main

  mathDR

  authored
  1c8f9de

Commits on Oct 18, 2024

• implementing pr feedback

  mathDR

  committed
  c252b51

Commits on Oct 21, 2024

• updated ademamix with author docstrings

  mathDR

  committed
  f9b6559

• added docstrings and matched adamw api

  mathDR

  committed
  4b621aa

• removed unneeded alpha scheduler

  mathDR

  committed
  1a51332

• added alpha as a scheduler

  mathDR

  committed
  4eb6065

• removed b3_scheduler

  mathDR

  committed
  4eb618a

• removed b3_scheduler

  mathDR

  committed
  9442085

• fixing tests with new docstrings

  mathDR

  committed
  420771f

• fixed docstring

  mathDR

  committed
  9690311

Commits on Oct 22, 2024

• updated notebook

  mathDR

  committed
  f892e29

• fixed import ordering

  mathDR

  committed
  81da0e4

Commits on Oct 24, 2024

• updated references using rst format

  mathDR

  committed
  94e3f0a

• updated docstrings

  mathDR

  committed
  40c0e6e

• fixed linting

  mathDR

  committed
  fb095e1

• synced ademamix api to adamw

  mathDR

  committed
  af8f22d

• added defaults to scale_by_ademamix

  mathDR

  committed
  47e4248

• fixed syntaxerror

  mathDR

  committed
  934761a

• updated docstrings

  mathDR

  committed
  1b72988

• fixed typo

  mathDR

  committed
  3e0699b

• reformatting note

  mathDR

  committed
  3ff8aba

• fixing formatting issues

  mathDR

  committed
  c933d17

• reformatting note

  mathDR

  committed
  bdeb3e2

• reformatting note

  mathDR

  committed
  8d38058

• fixed ademamix docstring

  mathDR

  committed
  6f4ec8a

• fixed notebook ordering and line lengths

  mathDR

  committed
  e8ba763

Commits on Oct 30, 2024

• added docs image for ademamix

  mathDR

  committed
  c84ce49

• added ademamix example to gallery

  mathDR

  committed
  dac4260

• reran notebook with colab link

  mathDR

  committed
  7c57abd