burn_ndarray/
backend.rs

1use crate::rand::NdArrayRng;
2use crate::{NdArrayQTensor, NdArrayTensor};
3use alloc::string::String;
4use burn_backend::quantization::{QuantMode, QuantScheme, QuantStore, QuantValue, quantizable};
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::sync::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();
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                        mode: QuantMode::Symmetric,
102                        #[cfg(not(feature = "export_tests"))]
103                            value: QuantValue::Q8F | QuantValue::Q8S,
104                        // For tests, "native" sub-byte quant serves as a reference for value equality.
105                        // Values are stored as i8 regardless.
106                        #[cfg(feature = "export_tests")]
107                            value:
108                            QuantValue::Q8F
109                            | QuantValue::Q8S
110                            | QuantValue::Q4F
111                            | QuantValue::Q4S
112                            | QuantValue::Q2F
113                            | QuantValue::Q2S,
114                        store: QuantStore::Native,
115                        ..
116                    // The value and store alone do not say which levels `quantize` will accept.
117                    } if quantizable(&scheme) => 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
212    /// A scheme this claims and then panics on is worse than one it declines, because the panic
213    /// lands on the first tensor rather than where the scheme was chosen.
214    #[test]
215    fn should_support_the_two_level_schemes_it_can_quantize() {
216        use burn_backend::ops::{FloatTensorOps, QTensorOps};
217        use burn_std::{ScaleDtype, TensorData};
218
219        type B = NdArray;
220        let device = NdArrayDevice::Cpu;
221        let scheme = device
222            .defaults()
223            .quantization
224            .scheme
225            .with_value(QuantValue::Q8S)
226            .per_block([4], ScaleDtype::UE4M3)
227            .per_tensor(ScaleDtype::F32);
228
229        assert!(B::supports_dtype(&device, DType::QFloat(scheme)));
230
231        let tensor = B::float_from_data(
232            TensorData::from([0.1f32, -0.4, 0.2, 0.9, -1.5, 0.3, 0.05, -0.02]),
233            &device,
234        );
235        let quantized = B::quantize_dynamic(tensor, &scheme);
236        assert_eq!(quantized.scheme.num_levels(), 2);
237
238        // Block scales already spanning f32's range are rejected by `quantize`.
239        assert!(!B::supports_dtype(
240            &device,
241            DType::QFloat(scheme.per_block([4], ScaleDtype::F32))
242        ));
243    }
244}
burn_ndarray/backend.rs

burn_ndarray/
backend.rs
