luma_tensor/grad/
store.rs

1use crate::{Device, Float, Tensor, TensorId, no_grad};
2use std::collections::HashMap;
3use std::ops::Index;
4
5/// Maps tensor ids to their accumulated gradients during a backward pass.
6/// Gradients are always `Float`-kind tensors on the same device.
7pub struct GradStore<D: Device>(HashMap<TensorId, Tensor<D, Float>>);
8
9impl<D: Device> GradStore<D> {
10    pub fn new() -> Self {
11        GradStore(HashMap::new())
12    }
13
14    pub fn get(&self, tensor: &Tensor<D, Float>) -> Option<&Tensor<D, Float>> {
15        self.0.get(&tensor.id())
16    }
17
18    pub fn get_by_id(&self, id: TensorId) -> Option<&Tensor<D, Float>> {
19        self.0.get(&id)
20    }
21
22    pub fn remove(&mut self, tensor: &Tensor<D, Float>) -> Option<Tensor<D, Float>> {
23        self.0.remove(&tensor.id())
24    }
25
26    pub fn insert(&mut self, tensor: &Tensor<D, Float>, grad: Tensor<D, Float>) -> Option<Tensor<D, Float>> {
27        self.0.insert(tensor.id(), grad)
28    }
29
30    /// Get the gradient accumulator for `tensor`, inserting a zeros tensor of the
31    /// same shape/dtype if absent.
32    pub fn or_insert(&mut self, tensor: &Tensor<D, Float>) -> crate::Result<&mut Tensor<D, Float>> {
33        use std::collections::hash_map::Entry;
34        let grad = match self.0.entry(tensor.id()) {
35            Entry::Occupied(e) => e.into_mut(),
36            Entry::Vacant(e) => e.insert(tensor.zeros_like()?),
37        };
38        Ok(grad)
39    }
40
41    pub fn get_ids(&self) -> impl Iterator<Item = &TensorId> {
42        self.0.keys()
43    }
44
45    pub fn tensors(&self) -> impl Iterator<Item = &Tensor<D, Float>> {
46        self.0.values()
47    }
48
49    pub fn iter(&self) -> std::collections::hash_map::Iter<'_, TensorId, Tensor<D, Float>> {
50        self.0.iter()
51    }
52
53    pub fn len(&self) -> usize {
54        self.0.len()
55    }
56
57    pub fn is_empty(&self) -> bool {
58        self.0.is_empty()
59    }
60
61    /// Global L2 norm of all gradients: `sqrt(sum ||g||^2)`.
62    pub fn global_norm(&self) -> crate::Result<f64> {
63        let mut total = 0.0f64;
64        for g in self.0.values() {
65            total += g.sqr()?.sum_all()?.to_scalar()?;
66        }
67        Ok(total.sqrt())
68    }
69
70    /// Clip all gradients by their global norm: if `norm > max_norm`, scale
71    /// every gradient by `max_norm / norm`. Returns the (pre-clip) norm.
72    pub fn clip_grad_norm(&mut self, max_norm: f64) -> crate::Result<f64> {
73        no_grad!();
74        let norm = self.global_norm()?;
75        if norm > max_norm {
76            let scale = max_norm / norm;
77            for g in self.0.values_mut() {
78                g.mul_scalar_(scale)?;
79            }
80        }
81        Ok(norm)
82    }
83
84    pub fn clear(&mut self) {
85        self.0.clear();
86    }
87}
88
89impl<D: Device> Default for GradStore<D> {
90    fn default() -> Self {
91        Self::new()
92    }
93}
94
95impl<D: Device> Index<&Tensor<D, Float>> for GradStore<D> {
96    type Output = Tensor<D, Float>;
97    fn index(&self, index: &Tensor<D, Float>) -> &Self::Output {
98        self.get(index).unwrap()
99    }
100}
luma_tensor/grad/store.rs

luma_tensor/grad/
store.rs
