burn_nn/modules/cache/base.rs
1use burn_core as burn;
2
3use burn::tensor::Tensor;
4use burn::tensor::backend::Backend;
5
6pub(crate) enum CacheState<T> {
7 Value(T),
8 Empty,
9}
10
11/// A cache for a tensor.
12pub struct TensorCache<B: Backend, const D: usize> {
13 pub(crate) state: CacheState<Tensor<B, D>>,
14}
15
16impl<B: Backend, const D: usize> TensorCache<B, D> {
17 /// Creates a new empty cache.
18 ///
19 /// # Returns
20 ///
21 /// The empty cache.
22 pub fn empty() -> Self {
23 Self {
24 state: CacheState::Empty,
25 }
26 }
27}