Struct LayoutOptimizedAdam 

pub struct LayoutOptimizedAdam { /* private fields */ }

Memory-optimized Adam optimizer using SoA layout.

Implementations

impl LayoutOptimizedAdam

pub fn new(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self

Creates a new layout-optimized Adam optimizer.

pub fn with_alignment( lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32, alignment: AlignmentConfig, ) -> Self

Creates an optimizer with custom alignment configuration.

pub fn avx512_optimized( lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32, ) -> Self

Creates an AVX-512 optimized variant.

pub fn layout_stats(&self) -> LayoutStats

Gets layout optimization statistics.

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

Adds a parameter to the optimizer with optimal layout.

Trait Implementations

impl Debug for LayoutOptimizedAdam

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

Formats the value using the given formatter.

impl Optimizer for LayoutOptimizedAdam

fn update(&mut self, parameter: &mut Tensor, grad: &Tensor) -> Result<()>

Updates a parameter based on its gradient.

fn zero_grad(&mut self)

Clears all accumulated gradients.

fn step(&mut self)

Performs a single optimization step.

fn get_lr(&self) -> f32

Gets the current learning rate.

fn set_lr(&mut self, lr: f32)

Sets a new learning rate.

fn accumulate_grad( &mut self, parameter: &mut Tensor, grad: &Tensor, ) -> Result<(), TrustformersError>

Accumulates gradients for gradient accumulation.

fn apply_accumulated_grads( &mut self, accumulation_steps: usize, ) -> Result<(), TrustformersError>

Applies accumulated gradients after gradient accumulation.

impl Send for LayoutOptimizedAdam

impl Sync for LayoutOptimizedAdam