Struct ModuleOptimizer 

pub struct ModuleOptimizer { /* private fields */ }

Optimizes a whole module by applying a per-parameter Optimizer to each of its parameters.

It is non-generic over the module and optimizer: any O: Optimizer is type-erased behind a dynamic optimizer, and per-parameter states are kept as type-erased states keyed by ParamId. Build one with optimizer.into() or OptimizerConfig::init().

It is possible to use different optimizers for different parameters. To do so, use the ModuleOptimizer::with_group function to add an optimizer for all parameters matching the provided group.

Implementations

impl ModuleOptimizer

pub fn has_gradient_clipping(&self) -> bool

Check if the optimizer has gradient clipping. If there are multiple optimizers, checks if any group has gradient clipping.

pub fn grad_clipping(&self) -> Option<&GradientClipping>

Access the gradient clipping. If there are multiple optimizers, returns the first optimizer’s GradientClipping.

pub fn with_grad_clipping( self, gradient_clipping: GradientClipping, ) -> ModuleOptimizer

Sets the gradient clipping. If there are multiple optimizers, assigns it to the first one.

Arguments

• gradient_clipping - The gradient clipping.

Returns

The optimizer.

pub fn with_group<O>( self, group: ParamGroup, optim: O, grad_clipping: Option<GradientClipping>, ) -> ModuleOptimizer

where O: DynOptimizer + 'static,

Adds an optimizer specific to a parameter group.

Parameters matching this group will be optimized using the provided optimizer and gradient clipping configuration.

Matching Rules

• Precedence: If a parameter matches multiple groups, the last group added takes precedence.
• Fallback: The first optimizer added must match all parameters to act as a global fallback.

Side Effects

• State Reset: Adding a new group will reset any existing optimizer states for parameters that match the new group.

impl ModuleOptimizer

pub fn step<M>( &mut self, lr_module: impl Into<ModuleLearningRate>, module: M, grads: GradientsParams, ) -> M

where M: Module,

Update the module parameters with the given gradients, advancing the optimizer state.

pub fn step_multi<M>( &mut self, lr_module: impl Into<ModuleLearningRate>, module: M, grads: MultiGradientsParams, ) -> M

where M: Module,

Like step, but accumulating gradients sourced from multiple devices.

pub fn to_record(&self) -> OptimizerRecord

Decompose the optimizer state into a serializable OptimizerRecord.

pub fn load_record(self, record: OptimizerRecord) -> ModuleOptimizer

Load the optimizer state from an OptimizerRecord.

State tensors are materialized on the default device; no device argument is needed because each parameter’s state is migrated to that parameter’s (gradient’s) device on the next step — see the to_device call in the step path. The load device is therefore irrelevant to correctness.

pub fn into_bytes(&self) -> Result<Bytes, RecordError>

Serialize the optimizer state to an in-memory burnpack byte buffer.

pub fn from_bytes(self, bytes: Bytes) -> Result<ModuleOptimizer, RecordError>

Load the optimizer state from an in-memory burnpack byte buffer.

pub fn save<P>(&self, path: P) -> Result<(), RecordError>

where P: AsRef<Path>,

Save the optimizer state to a burnpack file on disk.

pub fn load<P>(self, path: P) -> Result<ModuleOptimizer, RecordError>

where P: AsRef<Path>,

Load the optimizer state from a burnpack file on disk.

Trait Implementations

impl Clone for ModuleOptimizer

fn clone(&self) -> ModuleOptimizer

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<O> From<O> for ModuleOptimizer

where O: Optimizer,

fn from(optim: O) -> ModuleOptimizer

Converts to this type from the input type.