Feature request for the AdeMAMix optimizer.

[mathDR]

The AdeMAMix optimizer is a simple modification of
the Adam optimizer with a mixture of two EMAs to better take advantage of past gradients.

The paper has optax skeleton code which I could contribute if the maintainers deem this a good fit for the repo.

[clementpoiret]

Hey, any update on it since the PR has been closed? :) thanks !

[mathDR]

Apologies. I closed this because It was just easier to "begin again". I have a draft PR locally that I was going to push this week.

[clementpoiret]

Awesome, good to know! Thx for your work 👌

[mathDR]

Okay PR here