Struct SharedGradientStore 

pub struct SharedGradientStore { /* private fields */ }

Thread-safe gradient store that simulates AllReduce across world_size ranks.

Each rank pushes its local gradients via SharedGradientStore::push. Once all ranks have pushed, any rank can call SharedGradientStore::all_reduce_mean to obtain the element-wise average. Call SharedGradientStore::clear after each optimiser step to reset state for the next iteration.

Implementations

impl SharedGradientStore

pub fn new(world_size: usize) -> Self

Create a new store for world_size ranks.

pub fn push(&self, rank: usize, grads: Vec<GradientBuffer>) -> ModelResult<()>

Submit gradient buffers from rank.

Errors

Returns an error if the mutex is poisoned or rank >= world_size.

pub fn all_reduce_mean(&self, _rank: usize) -> ModelResult<Vec<GradientBuffer>>

Wait until all ranks have pushed, then return the element-wise mean.

In tests all ranks run in the same process/thread, so all buffers will already be filled before this is called.

Errors

Returns an error if not all ranks have submitted yet, or on lock failure.

pub fn clear(&self) -> ModelResult<()>

Clear all gradient buffers — call after each optimiser step.

Errors

Returns an error on lock failure.