bimm-contracts 0.2.2

Runtime contracts for the bimm 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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

use burn::prelude::{Backend, Shape, Tensor};
use burn::tensor::BasicOps;

/// A trait that provides a method to extract the shape from various types.
pub trait ShapeArgument {
    /// Extracts the shape from the implementing type.
    fn get_shape(self) -> Shape;
}

impl ShapeArgument for &Shape {
    fn get_shape(self) -> Shape {
        self.clone()
    }
}

impl ShapeArgument for Shape {
    fn get_shape(self) -> Shape {
        self
    }
}

impl<const D: usize> ShapeArgument for &[usize; D] {
    fn get_shape(self) -> Shape {
        Shape::from(*self)
    }
}

impl<const D: usize> ShapeArgument for &[u32; D] {
    fn get_shape(self) -> Shape {
        Shape::from(self.map(|d| d as usize))
    }
}

impl<const D: usize> ShapeArgument for &[i32; D] {
    fn get_shape(self) -> Shape {
        Shape::from(self.map(|d| d as usize))
    }
}

impl ShapeArgument for &[usize] {
    fn get_shape(self) -> Shape {
        Shape::from(self.to_vec())
    }
}

impl ShapeArgument for &[u32] {
    fn get_shape(self) -> Shape {
        Shape::from(self.iter().map(|&d| d as usize).collect::<Vec<_>>())
    }
}

impl ShapeArgument for &[i32] {
    fn get_shape(self) -> Shape {
        Shape::from(self.iter().map(|&d| d as usize).collect::<Vec<_>>())
    }
}

impl ShapeArgument for &Vec<usize> {
    fn get_shape(self) -> Shape {
        Shape::from(self)
    }
}

impl ShapeArgument for &Vec<u32> {
    fn get_shape(self) -> Shape {
        Shape::from(self.iter().map(|&d| d as usize).collect::<Vec<_>>())
    }
}

impl<B, const D: usize, K> ShapeArgument for &Tensor<B, D, K>
where
    B: Backend,
    K: BasicOps<B>,
{
    fn get_shape(self) -> Shape {
        self.shape()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn::prelude::Shape;

    #[test]
    fn test_shape_argument() {
        let shape = Shape::from([2, 3, 4]);
        assert_eq!(shape.clone().get_shape(), shape);

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

        {
            let arr: [usize; 3] = [2, 3, 4];
            assert_eq!(arr.get_shape(), shape);

            let arr_ref: &[usize] = &arr;
            assert_eq!(arr_ref.get_shape(), shape);
        }

        {
            let arr: [u32; 3] = [2, 3, 4];
            assert_eq!(arr.get_shape(), shape);

            let arr_ref: &[u32] = &arr;
            assert_eq!(arr_ref.get_shape(), shape);
        }

        {
            let arr: [i32; 3] = [2, 3, 4];
            assert_eq!(arr.get_shape(), shape);

            let arr_ref: &[i32] = &arr;
            assert_eq!(arr_ref.get_shape(), shape);
        }

        {
            let vec: Vec<usize> = vec![2, 3, 4];
            assert_eq!(vec.get_shape(), shape);

            let vec_ref: &Vec<usize> = &vec;
            assert_eq!(vec_ref.get_shape(), shape);
        }

        {
            let vec: Vec<u32> = vec![2, 3, 4];
            assert_eq!(vec.get_shape(), shape);

            let vec_ref: &Vec<u32> = &vec;
            assert_eq!(vec_ref.get_shape(), shape);
        }

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