Struct GradStore 

pub struct GradStore<D: Device>(/* private fields */);

Maps tensor ids to their accumulated gradients during a backward pass. Gradients are always Float-kind tensors on the same device.

Implementations

impl<D: Device> GradStore<D>

pub fn new() -> Self

pub fn get(&self, tensor: &Tensor<D, Float>) -> Option<&Tensor<D, Float>>

pub fn get_by_id(&self, id: TensorId) -> Option<&Tensor<D, Float>>

pub fn remove(&mut self, tensor: &Tensor<D, Float>) -> Option<Tensor<D, Float>>

pub fn insert( &mut self, tensor: &Tensor<D, Float>, grad: Tensor<D, Float>, ) -> Option<Tensor<D, Float>>

pub fn or_insert( &mut self, tensor: &Tensor<D, Float>, ) -> Result<&mut Tensor<D, Float>>

Get the gradient accumulator for tensor, inserting a zeros tensor of the same shape/dtype if absent.

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

pub fn tensors(&self) -> impl Iterator<Item = &Tensor<D, Float>>

pub fn iter(&self) -> Iter<'_, TensorId, Tensor<D, Float>>

pub fn len(&self) -> usize

pub fn is_empty(&self) -> bool

pub fn global_norm(&self) -> Result<f64>

Global L2 norm of all gradients: sqrt(sum ||g||^2).

pub fn clip_grad_norm(&mut self, max_norm: f64) -> Result<f64>

Clip all gradients by their global norm: if norm > max_norm, scale every gradient by max_norm / norm. Returns the (pre-clip) norm.

pub fn clear(&mut self)

Trait Implementations

impl<D: Device> Default for GradStore<D>

fn default() -> Self

Returns the “default value” for a type.

impl<D: Device> Index<&Tensor<D>> for GradStore<D>

type Output = Tensor<D>

The returned type after indexing.

fn index(&self, index: &Tensor<D, Float>) -> &Self::Output

Performs the indexing (container[index]) operation.