bimm-contracts 0.20.1

Runtime shape/geometry contracts for the burn framework.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

//! Burn Framework Support

use crate::ShapeArgument;
use alloc::vec::Vec;
use burn::prelude::{Backend, Shape, Tensor};
use burn::tensor::{BasicOps, TensorKind};

impl ShapeArgument for &Shape {
    fn get_shape_vec(self) -> Vec<usize> {
        self.dims.clone()
    }
}

impl ShapeArgument for Shape {
    fn get_shape_vec(self) -> Vec<usize> {
        self.dims
    }
}

impl<B, const D: usize, K> ShapeArgument for &Tensor<B, D, K>
where
    B: Backend,
    K: TensorKind<B> + BasicOps<B>,
{
    fn get_shape_vec(self) -> Vec<usize> {
        self.dims().to_vec()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::vec;

    #[test]
    fn test_shape_argument() {
        let expected = vec![2, 3, 4];

        let shape = Shape::from([2, 3, 4]);
        assert_eq!(&shape.clone().get_shape_vec(), &expected);

        let shape_ref: &Shape = &shape;
        assert_eq!(&shape_ref.get_shape_vec(), &expected);

        let tensor: Tensor<burn::backend::NdArray, 2> = Tensor::zeros([2, 2], &Default::default());
        let tensor_ref = &tensor;
        assert_eq!(&tensor_ref.get_shape_vec(), &[2, 2]);
    }
}