burn_std/tensor/
container.rs

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

burn_std/tensor/
container.rs
