An implementation of the optimizer published in "Promoting Exploration in Memory-Augmented Adam using Critical Momenta" at https://arxiv.org/abs/2307.09638. I implemented this after a brief review of the paper, and while I believe it to be correct, I cannot guarantee it is.
Happy hacking!