Struct Muon 

pub struct Muon {
    pub lr: f32,
    pub momentum: f32,
    pub nesterov: bool,
    pub weight_decay: f32,
    pub ns_steps: u32,
    pub ns_coeffs: (f32, f32, f32),
    /* private fields */
}

Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. Muon — Momentum-Orthogonalized-by-Newton-Schulz.

Per-tensor state: one momentum buffer per matrix (half of Adam’s footprint, like Lion).

Fields

lr: f32

Learning rate. The Newton–Schulz update has roughly unit Frobenius norm per column, so this is on the same scale as SGD’s lr — typically 2e-2 to 5e-2.

momentum: f32

Polyak momentum coefficient. Default 0.95.

nesterov: bool

Use Nesterov lookahead inside the matrix being orthogonalized. Default true.

weight_decay: f32

Decoupled weight-decay coefficient λ. Default 0.0.

ns_steps: u32

Newton–Schulz iteration count. 5 is the published default; 3 is enough for most well-conditioned matrices.

ns_coeffs: (f32, f32, f32)

(a, b, c) coefficients of the cubic Newton–Schulz iteration X ← a·X + b·(XXᵀ)X + c·(XXᵀ)²X. Defaults match Jordan et al.

Implementations

impl Muon

pub fn new(lr: f32) -> Muon

Construct with (μ, nesterov, λ, ns_steps) = (0.95, true, 0.0, 5) and the published NS coefficients.

pub fn with_momentum(self, mu: f32) -> Muon

Override the Polyak momentum coefficient.

pub fn with_weight_decay(self, wd: f32) -> Muon

Override the decoupled-decay coefficient.

pub fn with_ns_steps(self, n: u32) -> Muon

Override the Newton–Schulz iteration count.

Trait Implementations

impl Clone for Muon

fn clone(&self) -> Muon

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Muon

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl Optimizer for Muon

fn set_lr(&mut self, lr: f32)

Set the base learning rate (for LR schedules / warmup). Default is a no-op for algorithms without a scalar lr (e.g. Adafactor’s relative step size); every algorithm in this crate that has an lr field overrides this to update it.

fn step(&mut self, name: &str, shape: &[usize], param: &mut [f32], grad: &[f32])

fn step_batch(&mut self, items: &mut [OptItem<'_>])

Batched step over ALL parameters in one call. Default: sequential step per item — bit-identical to the per-parameter loop. Optimizers whose parameter groups are independent (e.g. Muon on the 2-D weight matrices vs AdamW on the embeddings/biases/norms) can override this to run the groups on separate threads; because the groups touch disjoint parameters and disjoint optimizer state, the result is bit-for-bit the same as the serial loop — only the wall-clock (the CPU-side optimizer bubble) shrinks toward max(group_times) instead of their sum.

fn end_iteration(&mut self)

Advance the global step counter. Most algorithms increment per call to [step], so most implementations leave this a no-op.

fn lr_scale(&self, _name: &str) -> f32

Per-tensor multiplier on the effective learning rate. Default is 1.0 for every name. Override when wrapping this crate to support per-name LR schedules (e.g. embedding-vs-attention splits, or the Gaussian-splat attribute-typed LR setup). The CPU impls in this crate currently honor this only when the caller passes a pre-scaled lr for the relevant call — backends are encouraged to consult it inside their fused kernel.