burn-mlx 0.1.2

MLX backend for Burn deep learning framework - native Apple Silicon GPU acceleration
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
138
139
140
141
142
143
144

//! Element type mappings between Burn and MLX.

use burn_tensor::{DType, Element};
use half::{bf16, f16};
use mlx_rs::Dtype;

/// Trait for elements that can be used with MLX.
pub trait MlxElement: Element + Clone + Send + Sync + 'static {
    /// Get the MLX data type for this element.
    fn mlx_dtype() -> Dtype;

    /// Get the Burn DType for this element.
    fn dtype() -> DType;
}

impl MlxElement for f32 {
    fn mlx_dtype() -> Dtype {
        Dtype::Float32
    }
    fn dtype() -> DType {
        DType::F32
    }
}

impl MlxElement for f64 {
    fn mlx_dtype() -> Dtype {
        Dtype::Float64
    }
    fn dtype() -> DType {
        DType::F64
    }
}

impl MlxElement for f16 {
    fn mlx_dtype() -> Dtype {
        Dtype::Float16
    }
    fn dtype() -> DType {
        DType::F16
    }
}

impl MlxElement for bf16 {
    fn mlx_dtype() -> Dtype {
        Dtype::Bfloat16
    }
    fn dtype() -> DType {
        DType::BF16
    }
}

impl MlxElement for i32 {
    fn mlx_dtype() -> Dtype {
        Dtype::Int32
    }
    fn dtype() -> DType {
        DType::I32
    }
}

impl MlxElement for i64 {
    fn mlx_dtype() -> Dtype {
        Dtype::Int64
    }
    fn dtype() -> DType {
        DType::I64
    }
}

impl MlxElement for i16 {
    fn mlx_dtype() -> Dtype {
        Dtype::Int16
    }
    fn dtype() -> DType {
        DType::I16
    }
}

impl MlxElement for i8 {
    fn mlx_dtype() -> Dtype {
        Dtype::Int8
    }
    fn dtype() -> DType {
        DType::I8
    }
}

impl MlxElement for u8 {
    fn mlx_dtype() -> Dtype {
        Dtype::Uint8
    }
    fn dtype() -> DType {
        DType::U8
    }
}

impl MlxElement for u16 {
    fn mlx_dtype() -> Dtype {
        Dtype::Uint16
    }
    fn dtype() -> DType {
        DType::U16
    }
}

impl MlxElement for u32 {
    fn mlx_dtype() -> Dtype {
        Dtype::Uint32
    }
    fn dtype() -> DType {
        DType::U32
    }
}

impl MlxElement for u64 {
    fn mlx_dtype() -> Dtype {
        Dtype::Uint64
    }
    fn dtype() -> DType {
        DType::U64
    }
}

impl MlxElement for bool {
    fn mlx_dtype() -> Dtype {
        Dtype::Bool
    }
    fn dtype() -> DType {
        DType::Bool
    }
}

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

    #[test]
    fn test_dtype_mapping() {
        assert_eq!(f32::mlx_dtype(), Dtype::Float32);
        assert_eq!(f64::mlx_dtype(), Dtype::Float64);
        assert_eq!(i32::mlx_dtype(), Dtype::Int32);
        assert_eq!(bool::mlx_dtype(), Dtype::Bool);
    }
}