burn_std/tensor/
container.rs

1use alloc::boxed::Box;
2use core::any::Any;
3
4#[cfg(not(feature = "std"))]
5use alloc::vec::Vec;
6#[cfg(not(feature = "std"))]
7use hashbrown::HashMap;
8
9#[cfg(feature = "std")]
10use std::collections::HashMap;
11
12/// Contains tensor of arbitrary dimension.
13#[derive(Debug)]
14pub struct TensorContainer<ID> {
15    tensors: HashMap<ID, Box<dyn Any + Send>>,
16}
17
18impl<ID> Default for TensorContainer<ID>
19where
20    ID: core::hash::Hash + PartialEq + Eq + core::fmt::Debug,
21{
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl<ID> TensorContainer<ID>
28where
29    ID: core::hash::Hash + PartialEq + Eq + core::fmt::Debug,
30{
31    /// Create an empty container.
32    pub fn new() -> Self {
33        Self {
34            tensors: HashMap::new(),
35        }
36    }
37
38    /// Get a tensor with the given ID.
39    pub fn get<T: Clone + Send + 'static>(&self, id: &ID) -> Option<T> {
40        let grad = self.tensors.get(id)?;
41
42        let tensor = grad.downcast_ref::<T>().unwrap();
43
44        Some(tensor.clone())
45    }
46
47    /// Get a mutable reference to the tensor with the given ID.
48    pub fn get_mut_ref<T: Clone + Send + 'static>(&mut self, id: &ID) -> Option<&mut T> {
49        let grad = self.tensors.get_mut(id)?;
50
51        let tensor = grad.downcast_mut::<T>().unwrap();
52
53        Some(tensor)
54    }
55
56    /// Register a new tensor for the given ID.
57    ///
58    /// # Notes
59    ///
60    /// If a tensor is already registered for the given ID, it will be replaced.
61    pub fn register<T: Clone + Send + 'static>(&mut self, id: ID, value: T) {
62        self.tensors.insert(id, Box::new(value));
63    }
64
65    /// Remove a tensor for the given ID and returns it.
66    pub fn remove<T: Clone + Send + 'static>(&mut self, id: &ID) -> Option<T> {
67        self.tensors
68            .remove(id)
69            .map(|item| *item.downcast::<T>().unwrap())
70    }
71
72    /// The number of tensors registered.
73    pub fn len(&self) -> usize {
74        self.tensors.len()
75    }
76
77    /// If any tensor is contained.
78    pub fn is_empty(&self) -> bool {
79        self.len() == 0
80    }
81
82    /// Get id of every tensor in the container
83    pub fn ids(&self) -> Vec<&ID> {
84        self.tensors.keys().collect()
85    }
86}
burn_std/tensor/container.rs

burn_std/tensor/
container.rs
