burn_ndarray/
backend.rs

1use crate::rand::NdArrayRng;
2use crate::{NdArrayQTensor, NdArrayTensor};
3use alloc::string::String;
4use burn_backend::quantization::{QuantLevel, QuantMode, QuantScheme, QuantStore, QuantValue};
5use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor};
6use burn_backend::{Backend, BackendTypes, DType, DeviceId, DeviceOps};
7use burn_ir::{BackendIr, HandleKind, TensorHandle};
8use burn_std::stub::Mutex;
9use burn_std::{BoolStore, DeviceSettings, QuantConfig};
10use rand::SeedableRng;
11
12pub(crate) static SEED: Mutex<Option<NdArrayRng>> = Mutex::new(None);
13
14/// The device type for the ndarray backend.
15#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
16pub enum NdArrayDevice {
17    /// The CPU device.
18    #[default]
19    Cpu,
20}
21
22impl DeviceOps for NdArrayDevice {
23    fn defaults(&self) -> DeviceSettings {
24        // E = f32, I = i64
25        DeviceSettings::new(
26            DType::F32,
27            DType::I64,
28            DType::Bool(BoolStore::Native),
29            QuantConfig::new(
30                QuantScheme::default().with_store(QuantStore::Native),
31                Default::default(),
32            ),
33        )
34    }
35}
36
37impl burn_backend::Device for NdArrayDevice {
38    fn from_id(_device_id: DeviceId) -> Self {
39        Self::Cpu
40    }
41
42    fn to_id(&self) -> DeviceId {
43        DeviceId {
44            type_id: 0,
45            index_id: 0,
46        }
47    }
48}
49
50/// Tensor backend that uses the [ndarray](ndarray) crate for executing tensor operations.
51///
52/// This backend is compatible with CPUs and can be compiled for almost any platform, including
53/// `wasm`, `arm`, and `x86`.
54#[derive(Clone, Copy, Default, Debug)]
55pub struct NdArray;
56
57impl BackendTypes for NdArray {
58    type Device = NdArrayDevice;
59
60    type FloatTensorPrimitive = NdArrayTensor;
61    type IntTensorPrimitive = NdArrayTensor;
62    type BoolTensorPrimitive = NdArrayTensor;
63    type QuantizedTensorPrimitive = NdArrayQTensor;
64
65    type GraphPrimitive = burn_backend::GraphUnsupported;
66}
67
68impl Backend for NdArray {
69    fn ad_enabled(_device: &Self::Device) -> bool {
70        false
71    }
72
73    fn name(_device: &Self::Device) -> String {
74        String::from("ndarray")
75    }
76
77    fn seed(_device: &Self::Device, seed: u64) {
78        let rng = NdArrayRng::seed_from_u64(seed);
79        let mut seed = SEED.lock().unwrap();
80        *seed = Some(rng);
81    }
82
83    fn dtype_usage(_device: &Self::Device, dtype: DType) -> burn_backend::DTypeUsageSet {
84        match dtype {
85            DType::F64
86            | DType::F32
87            | DType::Flex32
88            | DType::I64
89            | DType::I32
90            | DType::I16
91            | DType::I8
92            | DType::U64
93            | DType::U32
94            | DType::U16
95            | DType::U8
96            | DType::Bool(BoolStore::Native) => burn_backend::DTypeUsage::general(),
97            DType::F16 | DType::BF16 | DType::Bool(_) => burn_backend::DTypeUsageSet::empty(),
98            DType::QFloat(scheme) => {
99                match scheme {
100                    QuantScheme {
101                        level: QuantLevel::Tensor | QuantLevel::Block(_),
102                        mode: QuantMode::Symmetric,
103                        #[cfg(not(feature = "export_tests"))]
104                            value: QuantValue::Q8F | QuantValue::Q8S,
105                        // For tests, "native" sub-byte quant serves as a reference for value equality.
106                        // Values are stored as i8 regardless.
107                        #[cfg(feature = "export_tests")]
108                            value:
109                            QuantValue::Q8F
110                            | QuantValue::Q8S
111                            | QuantValue::Q4F
112                            | QuantValue::Q4S
113                            | QuantValue::Q2F
114                            | QuantValue::Q2S,
115                        store: QuantStore::Native,
116                        ..
117                    } => burn_backend::DTypeUsage::general(),
118                    _scheme => burn_backend::DTypeUsageSet::empty(),
119                }
120            }
121        }
122    }
123
124    fn device_count(_: u16) -> usize {
125        1
126    }
127
128    fn flush(_device: &Self::Device) {}
129}
130
131impl BackendIr for NdArray {
132    type Handle = HandleKind<Self>;
133
134    fn float_tensor(handle: TensorHandle<Self::Handle>) -> FloatTensor<Self> {
135        match handle.handle {
136            HandleKind::Float(handle) => handle,
137            _ => panic!("Expected float handle, got {}", handle.handle.name()),
138        }
139    }
140
141    fn int_tensor(handle: TensorHandle<Self::Handle>) -> IntTensor<Self> {
142        match handle.handle {
143            HandleKind::Int(handle) => handle,
144            _ => panic!("Expected int handle, got {}", handle.handle.name()),
145        }
146    }
147
148    fn bool_tensor(handle: TensorHandle<Self::Handle>) -> BoolTensor<Self> {
149        match handle.handle {
150            HandleKind::Bool(handle) => handle,
151            _ => panic!("Expected bool handle, got {}", handle.handle.name()),
152        }
153    }
154
155    fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> QuantizedTensor<Self> {
156        match handle.handle {
157            HandleKind::Quantized(handle) => handle,
158            _ => panic!("Expected quantized handle, got {}", handle.handle.name()),
159        }
160    }
161
162    fn float_tensor_handle(tensor: FloatTensor<Self>) -> Self::Handle {
163        HandleKind::Float(tensor)
164    }
165
166    fn int_tensor_handle(tensor: IntTensor<Self>) -> Self::Handle {
167        HandleKind::Int(tensor)
168    }
169
170    fn bool_tensor_handle(tensor: BoolTensor<Self>) -> Self::Handle {
171        HandleKind::Bool(tensor)
172    }
173
174    fn quantized_tensor_handle(tensor: QuantizedTensor<Self>) -> Self::Handle {
175        HandleKind::Quantized(tensor)
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn should_support_dtypes() {
185        type B = NdArray;
186        let device = NdArrayDevice::Cpu;
187        let scheme = device.defaults().quantization.scheme;
188
189        assert!(B::supports_dtype(&device, DType::F64));
190        assert!(B::supports_dtype(&device, DType::F32));
191        assert!(B::supports_dtype(&device, DType::Flex32));
192        assert!(B::supports_dtype(&device, DType::I64));
193        assert!(B::supports_dtype(&device, DType::I32));
194        assert!(B::supports_dtype(&device, DType::I16));
195        assert!(B::supports_dtype(&device, DType::I8));
196        assert!(B::supports_dtype(&device, DType::U64));
197        assert!(B::supports_dtype(&device, DType::U32));
198        assert!(B::supports_dtype(&device, DType::U16));
199        assert!(B::supports_dtype(&device, DType::U8));
200        assert!(B::supports_dtype(&device, DType::Bool(BoolStore::Native)));
201        assert!(B::supports_dtype(&device, DType::QFloat(scheme)));
202
203        assert!(!B::supports_dtype(&device, DType::F16));
204        assert!(!B::supports_dtype(&device, DType::BF16));
205        // QuantStore::U32 not supported
206        assert!(!B::supports_dtype(
207            &device,
208            DType::QFloat(QuantScheme::default())
209        ));
210    }
211}
burn_ndarray/backend.rs

burn_ndarray/
backend.rs
