Struct Muon 

pub struct Muon<B: Backend> { /* private fields */ }

Muon optimizer.

Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step using a finite quintic Newton-Schulz iteration. This is an approximate spectral transformation, NOT an exact polar decomposition or a guarantee that U*U^T is identity. This implementation computes in the tensor’s dtype; it does not silently cast to BF16 or create FP32 master parameters.

Important Notes

1. Only for nonempty 2D parameters: Muon is designed for hidden weight matrices. Use AdamW or SGD for biases, embeddings, and layer norms.

2. Learning rate adjustment: Muon automatically adjusts the learning rate based on parameter shape. See AdjustLrFn for details.

3. Weight decay timing: Unlike typical optimizers, Muon applies weight decay AFTER orthogonalization but uses the original (unadjusted) learning rate for it.

Implementations

impl<B: Backend> Muon<B>

pub fn validate_step<const D: usize>( &self, lr: LearningRate, tensor: &Tensor<B, D>, grad: &Tensor<B, D>, state: Option<&MuonState<B, D>>, ) -> Result<(), MuonError>

Check tensor metadata before submitting kernels. Does not scan values for NaN/Inf. Unscale and validate gradients before calling (on every rank).

pub fn try_step<const D: usize>( &self, lr: LearningRate, tensor: Tensor<B, D>, grad: Tensor<B, D>, state: Option<MuonState<B, D>>, ) -> Result<(Tensor<B, D>, Option<MuonState<B, D>>), MuonError>

Fallible metadata-checked update. Output and state may still be executing asynchronously. An accepted submission is not proof of device completion.

Uses a complete 2D matrix, never a flattened mixed-parameter buffer or an arbitrary shard. Missing gradients must be skipped by the caller.

Trait Implementations

impl<B: Clone + Backend> Clone for Muon<B>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<B: Backend> SimpleOptimizer<B> for Muon<B>

fn step<const D: usize>( &self, lr: LearningRate, tensor: Tensor<B, D>, grad: Tensor<B, D>, state: Option<Self::State<D>>, ) -> (Tensor<B, D>, Option<Self::State<D>>)

Perform a single Muon optimization step.

Algorithm

1. Apply momentum to gradient
2. Orthogonalize update via Newton-Schulz
3. Adjust learning rate based on parameter shape
4. Apply weight decay (using original lr)
5. Update parameter (using adjusted lr)

Notes

Unlike typical optimizers, the weight decay and parameter update use different learning rates:

• Weight decay uses the original lr
• Parameter update uses the shape-adjusted lr

Panics

This function will panic if the input tensors are not 2D.

type State<const D: usize> = MuonState<B, D>

The state of the optimizer. It also implements record, so that it can be saved.

fn to_device<const D: usize>( state: Self::State<D>, device: &Device<B>, ) -> Self::State<D>

Change the device of the state.