Custom Optimizers
=================

MuonWithAuxAdam
---------------

MuonWithAuxAdam is a custom optimizer that combines Muon updates with an auxiliary Adam optimizer. This Muon implementation is based on the original Keller Jordan GitHub repository: https://github.com/KellerJordan/Muon/blob/master/muon.py

.. autoclass:: nequip.train.MuonWithAuxAdam