bunsen 0.21.0

bunsen is acceleration tooling for burn
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

use std::{
    fmt::Debug,
    ops::Deref,
};

use burn::module::{
    Param,
    ParamId,
    Parameter,
};

/// This is meta-descriptor for a [`Param`].
///
/// This type acts as [`AsRef<T>`], [`Deref<T>`].
///
/// Currently, this will always be `Param<Tensor<_, _, _>>`.
#[derive(Debug, Clone, PartialEq)]
pub struct ParamDesc<T>
where
    T: Debug + Clone + Send + PartialEq,
{
    param_id: ParamId,
    data: T,
}

impl<T> Deref for ParamDesc<T>
where
    T: Debug + Clone + Send + PartialEq,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

impl<T> AsRef<T> for ParamDesc<T>
where
    T: Debug + Clone + Send + PartialEq,
{
    fn as_ref(&self) -> &T {
        &self.data
    }
}

impl<T> ParamDesc<T>
where
    T: Debug + Clone + Send + PartialEq,
{
    /// Create a new [`ParamDesc`].
    pub fn new(
        param_id: ParamId,
        param: T,
    ) -> Self {
        Self {
            param_id,
            data: param,
        }
    }

    /// Get the [`ParamId`].
    pub fn param_id(&self) -> ParamId {
        self.param_id
    }
}

impl<T, D> From<&Param<T>> for ParamDesc<D>
where
    T: Parameter,
    D: for<'a> From<&'a T> + Debug + Clone + Send + PartialEq + 'static,
{
    fn from(param: &Param<T>) -> Self {
        Self::new(param.id, D::from(&param.val()))
    }
}

#[cfg(test)]
#[allow(unused_imports)]
mod tests {
    use burn::{
        nn::LinearConfig,
        prelude::Shape,
        tensor::DType,
    };

    use super::{
        super::{
            TensorDesc,
            TensorKindDesc,
        },
        *,
    };
    use crate::support::testing::SetupTestBackend;

    type B = SetupTestBackend;

    #[test]
    fn test_from_param() {
        let device = Default::default();

        let linear = LinearConfig::new(2, 3).init::<B>(&device);

        let param = linear.weight;

        let param_desc: ParamDesc<TensorDesc> = (&param).into();

        assert_eq!(param_desc.param_id(), param.id);
        assert_eq!(param_desc.kind(), TensorKindDesc::Float);
        assert_eq!(param_desc.dtype(), DType::F32);
        assert_eq!(param_desc.shape(), &Shape::new([2, 3]));

        assert_eq!(param_desc.rank(), 2);
        assert_eq!(param_desc.size_estimate(), DType::F32.size() * 2 * 3);
    }
}