Feature muon

[NewBornRustacean]

Add muon optimizer to burn-optim

What's new?:

• Adds a new Muon optimizer implementation (crates/burn-optim/src/optim/muon.rs).
• Implements Newton–Schulz orthogonalization, momentum integration, optional weight decay, and shape-based learning-rate adjustment (Original and MatchRmsAdamW modes).
• Includes comprehensive unit tests.

Notes:

• The original implementation uses bfloat16 and so is pytorch implementation.
• for now in this pr, the module tested with f32. bfloat16 tests could be added soon.

Test summary

• test_adjust_lr_fn_original - Verifies the Original learning-rate adjustment ratios for square, tall, and wide matrices.
• test_adjust_lr_fn_match_rms_adamw - Verifies the MatchRmsAdamW learning-rate adjustment ratios for example shapes.
• test_1d_tensor_panics- Ensures Newton–Schulz orthogonalization panics for 1D tensors (requires 2D).
• test_muon_optimizer_save_load_state - Verifies optimizer state can be saved and loaded for a Linear layer without bias.
• test_muon_with_weight_decay - Ensures weight decay is applied (weights are reduced) for a Linear layer without bias.
• test_newton_schulz_orthogonalization - Checks Newton–Schulz produces approximately orthogonal output (A * A^T ≈ I).
• test_tall_matrix_transpose - Ensures tall matrices are transposed internally and shape is preserved; verifies orthogonalization changes values and wide-matrix behavior.
• test_zero_gradient — Confirms Muon handles zero gradients without NaNs, creates state, and weight decay still reduces values when gradients are zero.

Related issue, readings, etc.

• Add muon optimizer #3924
• Keller Jordan's writeup - Muon: An optimizer for hidden layers in neural networks
• Keller Jordan's implementation
• pytorch implementation _muon.py

[nathanielsimard]

Are those parameters ignored during training if you use only a single optimizer?

[nathanielsimard]

Unsure if this should be the default behavior. There isn't a great way yet to define multiple optimizers for a single burn module (e.g a linear layer with a bias vector). Do you have an idea: @laggui ?