Struct GradientAccumulator 

pub struct GradientAccumulator { /* private fields */ }

Gradient accumulator managing multiple parameter gradients.

Provides micro-batching support by accumulating gradients across multiple forward/backward passes before applying an optimizer step.

Implementations

impl GradientAccumulator

pub fn new(config: AccumulationConfig) -> Self

Create a new gradient accumulator with the given configuration.

pub fn register(&mut self, name: impl Into<String>, shape: Vec<usize>)

Register a parameter with its gradient shape.

If the parameter is already registered, this is a no-op.

pub fn accumulate( &mut self, name: &str, grad: &[f64], ) -> Result<(), AccumulationError>

Accumulate a gradient for a named parameter.

Returns an error if the parameter has not been registered or if the gradient size does not match the registered shape.

pub fn should_update(&self) -> bool

Check if enough micro-batches have been accumulated to trigger an update.

pub fn get_gradients( &self, ) -> Result<HashMap<String, Vec<f64>>, AccumulationError>

Get all accumulated gradients, optionally normalized and clipped.

pub fn reset(&mut self)

Reset all buffers after an update step.

pub fn step( &mut self, gradients: &HashMap<String, Vec<f64>>, ) -> Result<bool, AccumulationError>

Accumulate a full micro-batch of gradients, returning true if an update should now be applied.

pub fn stats(&self) -> AccumulationStats

Get statistics about the accumulation state.