Struct Soap 

pub struct Soap {
    pub lr: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub shampoo_beta: f32,
    pub eps: f32,
    pub weight_decay: f32,
    pub precond_freq: u64,
    pub jacobi_sweeps: u32,
    /* private fields */
}

SOAP — Shampoo-in-Adam-basis optimizer.

Fields

lr: f32

Learning rate.

beta1: f32

First-moment EMA decay (in the rotated basis). Default 0.95.

beta2: f32

Second-moment EMA decay (in the rotated basis). Default 0.95.

shampoo_beta: f32

Decay for the L/R covariance EMAs. Often equal to β₂.

eps: f32

Denominator stability constant. Default 1e-8.

weight_decay: f32

Decoupled weight-decay coefficient λ. Default 0.01.

precond_freq: u64

Recompute the eigenbasis every precond_freq steps. Larger values amortize the Jacobi cost but lag the preconditioner. Default 10.

jacobi_sweeps: u32

Max Jacobi sweeps per rediagonalization. Default 30.

Implementations

impl Soap

pub fn new(lr: f32) -> Self

Construct with (β₁, β₂, sb, ε, λ, freq, sweeps) = (0.95, 0.95, 0.95, 1e-8, 0.01, 10, 30).

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

Override the decoupled-decay coefficient.

Trait Implementations

impl Clone for Soap

fn clone(&self) -> Soap

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Soap

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

Formats the value using the given formatter.

impl Optimizer for Soap

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

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.