Struct SoAOptimizerState 

pub struct SoAOptimizerState { /* private fields */ }

Structure of Arrays (SoA) layout for optimizer state.

Momentum and variance for every registered parameter live in two contiguous arrays rather than in per-parameter allocations, so a sweep over one moment touches consecutive cache lines. Each parameter owns the half-open range [momentum_offset, momentum_offset + size) of SoAOptimizerState::momentum (and likewise for the variance array); ranges are padded up to a cache-line boundary so no two parameters share a cache line.

The moments are persistent: update_parameter_soa reads the previous values out of these arrays and writes the new ones back. An earlier revision recomputed both from the current gradient on every call, which silently removed all of Adam’s adaptivity.

Implementations

impl SoAOptimizerState

pub fn new(alignment: AlignmentConfig) -> Self

Creates a new SoA optimizer state.

pub fn add_parameter(&mut self, id: String, size: usize) -> Result<()>

Adds a parameter to the SoA layout, reserving zeroed moment storage for it.

Registering the same id twice is a no-op, so callers may register lazily.

Errors

Returns an error when size is zero, which would produce an empty block.

pub fn get_parameter_info(&self, id: &str) -> Option<&ParameterInfo>

Gets parameter information by ID.

pub fn momentum_of(&self, id: &str) -> Option<&[f32]>

Reads the persisted first moment of a parameter (primarily for tests).

pub fn variance_of(&self, id: &str) -> Option<&[f32]>

Reads the persisted second moment of a parameter (primarily for tests).

pub fn step_count(&self) -> usize

Current global step counter.

pub fn update_parameter_soa( &mut self, param_id: &str, param: &mut [f32], grad: &[f32], lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32, ) -> Result<()>

Updates momentum and variance for a parameter using optimized memory access.

pub fn layout_stats(&self) -> LayoutStats

Gets memory layout statistics.

Trait Implementations

impl Debug for SoAOptimizerState

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

Formats the value using the given formatter.

impl Send for SoAOptimizerState

impl Sync for SoAOptimizerState