Struct MuonAdamW 

pub struct MuonAdamW<M: AutodiffModule<B>, B: AutodiffBackend> { /* private fields */ }

Mixed optimizer with fixed, explicit parameter identities. Missing gradients skip both momentum and decay for that parameter. A tied parameter is updated once, by its ParamId.

Implementations

impl<M: AutodiffModule<B>, B: AutodiffBackend> MuonAdamW<M, B>

pub fn muon_parameter_count(&self) -> usize

Number of distinct parameters routed to Muon, not number of matrix elements.

pub fn try_step_with_lrs( &mut self, muon_lr: LearningRate, adamw_lr: LearningRate, module: M, grads: GradientsParams, ) -> Result<M, MuonError>

Update with independent learning rates. Metadata for both groups is checked before either group submits an update. Device runtime errors are still asynchronous and this method is NOT a two-phase device transaction.

pub fn try_step_or_skip( &mut self, muon_lr: LearningRate, adamw_lr: LearningRate, module: M, grads: GradientsParams, skip_update: bool, ) -> Result<M, MuonError>

An explicit all-group skip makes no optimizer update and changes no state. The caller must decide skip consistently across all replicas and must unscale/check finite gradients before a non-skipped update.

pub fn try_load_record( self, record: MuonAdamWRecord<B>, ) -> Result<Self, MuonError>

Restore only compatible grouping/configuration. Does not silently move a momentum buffer between SGD and EMA conventions or between parameter roles.

Trait Implementations

impl<M: Clone + AutodiffModule<B>, B: Clone + AutodiffBackend> Clone for MuonAdamW<M, B>

where B::InnerBackend: Clone,

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<M: AutodiffModule<B>, B: AutodiffBackend> Optimizer<M, B> for MuonAdamW<M, B>

type Record = MuonAdamWRecord<B>

Optimizer associative type to be used when saving and loading the state.

fn step(&mut self, lr: LearningRate, module: M, grads: GradientsParams) -> M

Perform the optimizer step using the given learning rate and gradients. The updated module is returned.

fn step_multi( &mut self, _lr: LearningRate, _module: M, _grads: MultiGradientsParams, ) -> M

Perform the optimizer step using the given learning rate and gradients. The updated module is returned.

fn to_record(&self) -> Self::Record

Get the current state of the optimizer as a record.

fn load_record(self, record: Self::Record) -> Self

Load the state of the optimizer as a record.