feat: add preconditioning and coef presets to muon

[massena-t]

Add Preconditioning to the Muon Optimizer

In this commit, we implement different preconditioning methods for the Muon optimizer. Notably:

• Frobenius norm: as previously done by default.
• Spectral norm: for more precise orthogonalization.
• Schatten-4 norm: based of the Gelfand formula, as described in https://arxiv.org/abs/2506.10935 (Section 3.4).
• AOL preconditioning: as described in the Turbo-Muon paper https://arxiv.org/abs/2512.04632.

These options allow the user to tweak the trade-off between convergence speed of the Newton-Schulz iterative scheme and the computational budget associated to the preconditioning method.

We also add options to choose the coefficients used in the Newton-Schulz iterations. Further development could implement the PolarExpress (https://arxiv.org/abs/2505.16932) or CAN methods (https://arxiv.org/abs/2506.10935).

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.