Struct GradStore 

pub struct GradStore<R: Runtime> { /* private fields */ }

Storage for gradients computed during backward pass

Gradients are stored by tensor ID and accumulated when a tensor is used multiple times in the computation graph.

Implementations

impl<R: Runtime> GradStore<R>

pub fn new() -> Self

Create a new empty gradient store

pub fn get(&self, id: TensorId) -> Option<&Tensor<R>>

Get the gradient for a tensor

pub fn insert(&mut self, id: TensorId, grad: Tensor<R>)

Insert a gradient (overwrites if exists)

pub fn contains(&self, id: TensorId) -> bool

Check if a gradient exists

pub fn remove(&mut self, id: TensorId) -> Option<Tensor<R>>

Remove and return a gradient

pub fn keys(&self) -> impl Iterator<Item = &TensorId>

Get all tensor IDs with gradients

pub fn len(&self) -> usize

Number of stored gradients

pub fn is_empty(&self) -> bool

Check if empty

pub fn clear(&mut self)

Clear all gradients

pub fn accumulate<F>(&mut self, id: TensorId, grad: Tensor<R>, add_fn: F)

where F: FnOnce(Tensor<R>, Tensor<R>) -> Tensor<R>,

Accumulate a gradient for a tensor

If no gradient exists for this tensor, stores the gradient. If a gradient already exists, adds the new gradient to the existing one.

Arguments

• id - The tensor ID to accumulate gradient for
• grad - The gradient tensor to accumulate
• add_fn - Function to add two tensors: fn(existing, new) -> sum

This is used when a tensor is used multiple times in the computation graph, requiring its gradients to be summed according to the chain rule.

pub fn try_accumulate<F>( &mut self, id: TensorId, grad: Tensor<R>, add_fn: F, ) -> Result<()>

where F: FnOnce(Tensor<R>, Tensor<R>) -> Result<Tensor<R>>,

Accumulate a gradient with a fallible addition function

Like accumulate, but the addition function can fail and return a Result. This is the preferred method for use in backward passes where tensor operations may fail.

pub fn accumulate_or_insert(&mut self, id: TensorId, grad: Tensor<R>)

Insert a gradient, overwriting any existing value

This is intentionally simpler than accumulate - it directly inserts the gradient without addition. Use this when you don’t have access to a TensorOps client, or when overwriting semantics are desired.

For proper gradient accumulation (adding to existing gradients), use accumulate with an add function instead.

Trait Implementations

impl<R: Runtime> Default for GradStore<R>

fn default() -> Self

Returns the “default value” for a type.