Struct GradientBucket 

pub struct GradientBucket { /* private fields */ }

A bucket for accumulating gradients before all-reduce.

Implementations

impl GradientBucket

pub fn new(capacity: usize) -> Self

Creates a new gradient bucket.

pub fn is_full(&self) -> bool

Checks if the bucket is full.

pub fn is_empty(&self) -> bool

Checks if the bucket is empty.

pub fn size(&self) -> usize

Returns the current size.

pub fn add(&mut self, tensor: &Tensor<f32>) -> bool

Adds a tensor to the bucket.

pub fn data(&self) -> &[f32]

Returns the flattened data.

pub fn data_mut(&mut self) -> &mut [f32]

Returns mutable flattened data.

pub fn clear(&mut self)

Clears the bucket.

pub fn extract(&self) -> Vec<Tensor<f32>>

Extracts tensors back from the bucket.