Add Polar Express Muon

[shermansiu]

According to the following paper, using Polar Express instead of Newton-Schultz in Muon leads to improvements in the validation loss.
https://arxiv.org/abs/2505.16932

It was also used in Karpathy's NanoChat repo to train a GPT-2 equivalent model for less than $100 (karpathy/nanochat#481)

It might be worthwhile to add it here.

Official implementation: https://github.com/NoahAmsel/PolarExpress/blob/main/polar_express.py

[shermansiu]

According to the reference implementation that was used in an optimizer (mentioned in the paper), it seems like we just need to set the Newton-Schultz function to the Polar Express function to add it here.

    def _initialize_polar_factorizer(self, polar_method, polar_args):
        """Initialize the polar factorization method based on the provided name and parameters."""
        if polar_method == "Keller":
            return zeropower_via_newtonschulz5  # Use the method directly
        elif polar_method == "Jiacheng":
            return jiacheng
        elif polar_method == "polarexpress":
            return PolarExpress 
        elif polar_method == "fast_polarexpress":
            return partial(FastApplyPolarExpress, restart_interval=3, shift_eps=1e-3)

PolarExpress is Algorithm 2 in the paper and FastApplyPolarExpress refers to Algorithm 4.

https://github.com/modichirag/GPT-opt/blob/70c80de690de7810c0a7d55d2a924c01f8e04ba6/gptopt/optim/muon.py#L184-L193

[shermansiu]

And for reference, here's the implementation in NanoChat: https://github.com/karpathy/nanochat/blob/16b8ac7da33010dc7efcd9292f895703f5cff33a/nanochat/optim.py#L110-L123