Add AdEMAMix Optimizer

[mathDR]

This PR adds the AdaMAMix optimizer from The arxiv preprint: THE ADEMAMIX OPTIMIZER:
BETTER, FASTER, OLDER

Closes #1058

The docs have been updated, along with the relevant files. Furthermore, I ran a similar "test" to replicate the Rosenbrock figure from the paper:

Currently the docstrings are implemented, but further descriptions should/could be added. I will reach out to the paper authors to assist with that (if they are willing).

[vroulet]

Thanks @mathDR for the contribution!
Optimizers that have not passed the test of time (too recent like this one) are put in the contrib folder (see https://optax.readthedocs.io/en/latest/development.html#inclusion-criteria).
Before continuing (making this optimizer in the contribs folder), can you wait for #1060 to be merged?
I would like all optimizers to abide to common tests to avoid issues like #1038.
PS: use an editor that enforces indents of two spaces not four.

[mathDR]

Okay thanks for the comments @vroulet. I will put an alert to check the status of #1060 and see when that is merged.

[vroulet]

Hello @mathDR
#1060 has been merged.
If you want you can continue this PR. Put the algorithm in the contrib folder and make sure it passes the common tests.

[mathDR]

I closed this and deleted the forked repo (will be easier to fork again and start anew)