Struct OptimizerState 

pub struct OptimizerState {
    pub step: usize,
    pub momentum: HashMap<String, Vec<f32>>,
    pub variance: HashMap<String, Vec<f32>>,
    pub third_moment: HashMap<String, Vec<f32>>,
    pub param_steps: HashMap<String, usize>,
    pub velocity: HashMap<String, Vec<f32>>,
    pub params: ParamRegistry,
}

Unified state management for optimizer parameters.

This struct provides a consistent interface for tracking optimizer state across different algorithms, reducing code duplication and memory overhead.

Fields

step: usize

Current step counter for bias correction and scheduling

momentum: HashMap<String, Vec<f32>>

First moment estimates (momentum buffers)

variance: HashMap<String, Vec<f32>>

Second moment estimates (squared gradient buffers)

third_moment: HashMap<String, Vec<f32>>

Optional third moment estimates (for higher-order methods)

param_steps: HashMap<String, usize>

Per-parameter step counts (for adaptive methods)

velocity: HashMap<String, Vec<f32>>

Velocity buffers for optimization methods like SGD with momentum

params: ParamRegistry

Stable parameter identity registry.

Every buffer map above is keyed by the canonical keys handed out by this registry. It is deliberately not serialised: identity is reconstructed on load from the checkpointed keys themselves (see crate::param_id), which keeps the on-disk state format unchanged.

Implementations

impl OptimizerState

pub fn new() -> Self

Creates a new optimizer state with empty buffers.

pub fn param_key_for_tensor(&mut self, tensor: &Tensor) -> Result<String>

Resolves the stable state key for an anonymous parameter tensor.

This is the supported replacement for the old format!("{:p}", …) idiom: the returned key survives checkpoint save/load, whereas a heap address does not.

Errors

Returns an error when the tensor dtype has no addressable buffer, or when the parameter ordering does not match a restored checkpoint.

pub fn param_key(&mut self, addr: usize, numel: usize) -> Result<String>

Resolves the stable state key for an anonymous parameter given its buffer address and element count.

Errors

See OptimizerState::param_key_for_tensor.

pub fn param_key_named( &mut self, name: &str, addr: usize, numel: usize, ) -> String

Resolves the stable state key for a parameter that has a caller-supplied name.

Named keys make checkpoint resume independent of parameter visit order and should be preferred wherever the surrounding API carries names.

pub fn restore_param_key(&mut self, key: &str, numel: usize) -> Result<ParamId>

Rebuilds one registry slot from a checkpointed state key.

Errors

Returns an error when key carries no recognised identity prefix.

pub fn get_or_create_momentum( &mut self, param_id: String, size: usize, ) -> &mut Vec<f32>

Gets or creates momentum buffer for a parameter.

pub fn get_or_create_variance( &mut self, param_id: String, size: usize, ) -> &mut Vec<f32>

Gets or creates variance buffer for a parameter.

pub fn get_or_create_third_moment( &mut self, param_id: String, size: usize, ) -> &mut Vec<f32>

Gets or creates third moment buffer for a parameter.

pub fn step(&mut self)

Increments the global step counter.

pub fn step_param(&mut self, param_id: String)

Increments the step counter for a specific parameter.

pub fn get_param_step(&self, param_id: &str) -> usize

Gets the step count for a specific parameter.

pub fn clear(&mut self)

Clears all state buffers to free memory.

The parameter identity registry is cleared alongside the buffers: leaving registrations behind while dropping their state would keep params.len() reporting parameters whose buffers no longer exist, and would leave the registry’s bind cursor advanced past slots that are ready to be reused.

pub fn memory_usage(&self) -> StateMemoryStats

Gets memory usage statistics.

Trait Implementations

impl Clone for OptimizerState

fn clone(&self) -> OptimizerState

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for OptimizerState

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

Formats the value using the given formatter.

impl Default for OptimizerState

fn default() -> Self

Returns the “default value” for a type.

impl<'de> Deserialize<'de> for OptimizerState

fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl Serialize for OptimizerState

fn serialize<__S>(&self, __serializer: __S) -> Result<__S::Ok, __S::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.