burn-tch 0.21.0

LibTorch backend for the Burn framework using the tch bindings.
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

use burn_backend::Element;
use burn_backend::{bf16, f16};

/// The element type for the tch backend.
pub trait TchElement: Element + tch::kind::Element {
    /// Returns the associated tensor kind for [`tch::kind::Element`].
    fn kind() -> tch::Kind {
        Self::KIND
    }
}

impl TchElement for f64 {}
impl TchElement for f32 {}
impl TchElement for f16 {}
impl TchElement for bf16 {
    fn kind() -> tch::Kind {
        let mut kind = <Self as tch::kind::Element>::KIND;
        // Incorrect kind mapping in tch definitions, force bfloat16
        if matches!(Self::dtype(), burn_backend::DType::BF16) && kind == tch::Kind::Half {
            kind = tch::Kind::BFloat16
        }
        kind
    }
}

impl TchElement for i64 {}
impl TchElement for i32 {}
impl TchElement for i16 {}
impl TchElement for i8 {}

impl TchElement for u8 {}

impl TchElement for bool {}

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

    #[test]
    fn test_elem_kinds() {
        assert_eq!(f64::kind(), tch::Kind::Double);
        assert_eq!(f32::kind(), tch::Kind::Float);
        assert_eq!(f16::kind(), tch::Kind::Half);
        assert_eq!(bf16::kind(), tch::Kind::BFloat16);
        assert_eq!(i64::kind(), tch::Kind::Int64);
        assert_eq!(i32::kind(), tch::Kind::Int);
        assert_eq!(i16::kind(), tch::Kind::Int16);
        assert_eq!(i8::kind(), tch::Kind::Int8);
        assert_eq!(bool::kind(), tch::Kind::Bool);
    }
}