burn_std/data/
compare.rs

1use alloc::format;
2use alloc::string::String;
3use num_traits::{Float, ToPrimitive};
4
5use super::TensorData;
6use crate::{BoolStore, DType, Element, ElementOrdered, bf16, f16};
7
8/// The tolerance used to compare to floating point numbers.
9///
10/// Generally, two numbers `x` and `y` are approximately equal if
11///
12/// ```text
13/// |x - y| < max(R * (|x + y|), A)
14/// ```
15///
16/// where `R` is the relative tolerance and `A` is the absolute tolerance.
17///
18///
19/// The most common way to initialize this struct is to use `Tolerance::<F>::default()`.
20/// In that case, the relative and absolute tolerances are computed using an heuristic based
21/// on the EPSILON and MIN_POSITIVE values of the given floating point type `F`.
22///
23/// Another common initialization is `Tolerance::<F>::rel_abs(1e-4, 1e-5).set_half_precision_relative(1e-2)`.
24/// This will use a sane default to manage values too close to 0.0 and
25/// use different relative tolerances depending on the floating point precision.
26#[derive(Debug, Clone, Copy)]
27pub struct Tolerance<F> {
28    relative: F,
29    absolute: F,
30}
31
32impl<F: Float> Default for Tolerance<F> {
33    fn default() -> Self {
34        Self::balanced()
35    }
36}
37
38impl<F: Float> Tolerance<F> {
39    /// Create a tolerance with strict precision setting.
40    pub fn strict() -> Self {
41        Self {
42            relative: F::from(0.00).unwrap(),
43            absolute: F::from(64).unwrap() * F::min_positive_value(),
44        }
45    }
46    /// Create a tolerance with balanced precision setting.
47    pub fn balanced() -> Self {
48        Self {
49            relative: F::from(0.005).unwrap(), // 0.5%
50            absolute: F::from(1e-5).unwrap(),
51        }
52    }
53
54    /// Create a tolerance with permissive precision setting.
55    pub fn permissive() -> Self {
56        Self {
57            relative: F::from(0.01).unwrap(), // 1.0%
58            absolute: F::from(0.01).unwrap(),
59        }
60    }
61    /// When comparing two numbers, this uses both the relative and absolute differences.
62    ///
63    /// That is, `x` and `y` are approximately equal if
64    ///
65    /// ```text
66    /// |x - y| < max(R * (|x + y|), A)
67    /// ```
68    ///
69    /// where `R` is the `relative` tolerance and `A` is the `absolute` tolerance.
70    pub fn rel_abs<FF: ToPrimitive>(relative: FF, absolute: FF) -> Self {
71        let relative = Self::check_relative(relative);
72        let absolute = Self::check_absolute(absolute);
73
74        Self { relative, absolute }
75    }
76
77    /// When comparing two numbers, this uses only the relative difference.
78    ///
79    /// That is, `x` and `y` are approximately equal if
80    ///
81    /// ```text
82    /// |x - y| < R * max(|x|, |y|)
83    /// ```
84    ///
85    /// where `R` is the relative `tolerance`.
86    pub fn relative<FF: ToPrimitive>(tolerance: FF) -> Self {
87        let relative = Self::check_relative(tolerance);
88
89        Self {
90            relative,
91            absolute: F::from(0.0).unwrap(),
92        }
93    }
94
95    /// When comparing two numbers, this uses only the absolute difference.
96    ///
97    /// That is, `x` and `y` are approximately equal if
98    ///
99    /// ```text
100    /// |x - y| < A
101    /// ```
102    ///
103    /// where `A` is the absolute `tolerance`.
104    pub fn absolute<FF: ToPrimitive>(tolerance: FF) -> Self {
105        let absolute = Self::check_absolute(tolerance);
106
107        Self {
108            relative: F::from(0.0).unwrap(),
109            absolute,
110        }
111    }
112
113    /// Change the relative tolerance to the given one.
114    pub fn set_relative<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
115        self.relative = Self::check_relative(tolerance);
116        self
117    }
118
119    /// Change the relative tolerance to the given one only if `F` is half precision.
120    pub fn set_half_precision_relative<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
121        if core::mem::size_of::<F>() == 2 {
122            self.relative = Self::check_relative(tolerance);
123        }
124        self
125    }
126
127    /// Change the relative tolerance to the given one only if `F` is single precision.
128    pub fn set_single_precision_relative<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
129        if core::mem::size_of::<F>() == 4 {
130            self.relative = Self::check_relative(tolerance);
131        }
132        self
133    }
134
135    /// Change the relative tolerance to the given one only if `F` is double precision.
136    pub fn set_double_precision_relative<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
137        if core::mem::size_of::<F>() == 8 {
138            self.relative = Self::check_relative(tolerance);
139        }
140        self
141    }
142
143    /// Change the absolute tolerance to the given one.
144    pub fn set_absolute<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
145        self.absolute = Self::check_absolute(tolerance);
146        self
147    }
148
149    /// Change the absolute tolerance to the given one only if `F` is half precision.
150    pub fn set_half_precision_absolute<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
151        if core::mem::size_of::<F>() == 2 {
152            self.absolute = Self::check_absolute(tolerance);
153        }
154        self
155    }
156
157    /// Change the absolute tolerance to the given one only if `F` is single precision.
158    pub fn set_single_precision_absolute<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
159        if core::mem::size_of::<F>() == 4 {
160            self.absolute = Self::check_absolute(tolerance);
161        }
162        self
163    }
164
165    /// Change the absolute tolerance to the given one only if `F` is double precision.
166    pub fn set_double_precision_absolute<FF: ToPrimitive>(mut self, tolerance: FF) -> Self {
167        if core::mem::size_of::<F>() == 8 {
168            self.absolute = Self::check_absolute(tolerance);
169        }
170        self
171    }
172
173    /// Checks if `x` and `y` are approximately equal given the tolerance.
174    pub fn approx_eq(&self, x: F, y: F) -> bool {
175        // See the accepted answer here
176        // https://stackoverflow.com/questions/4915462/how-should-i-do-floating-point-comparison
177
178        // This also handles the case where both a and b are infinity so that we don't need
179        // to manage it in the rest of the function.
180        if x == y {
181            return true;
182        }
183
184        let diff = (x - y).abs();
185        let max = F::max(x.abs(), y.abs());
186
187        diff < self.absolute.max(self.relative * max)
188    }
189
190    fn check_relative<FF: ToPrimitive>(tolerance: FF) -> F {
191        let tolerance = F::from(tolerance).unwrap();
192        assert!(tolerance <= F::one());
193        tolerance
194    }
195
196    fn check_absolute<FF: ToPrimitive>(tolerance: FF) -> F {
197        let tolerance = F::from(tolerance).unwrap();
198        assert!(tolerance >= F::zero());
199        tolerance
200    }
201}
202
203impl TensorData {
204    /// Asserts the data is equal to another data.
205    ///
206    /// # Arguments
207    ///
208    /// * `other` - The other data.
209    /// * `strict` - If true, the data types must the be same.
210    ///   Otherwise, the comparison is done in the current data type.
211    ///
212    /// # Panics
213    ///
214    /// Panics if the data is not equal.
215    #[track_caller]
216    pub fn assert_eq(&self, other: &Self, strict: bool) {
217        if strict {
218            assert_eq!(
219                self.dtype, other.dtype,
220                "Data types differ ({:?} != {:?})",
221                self.dtype, other.dtype
222            );
223        }
224
225        match self.dtype {
226            DType::F64 => self.assert_eq_elem::<f64>(other),
227            DType::F32 | DType::Flex32 => self.assert_eq_elem::<f32>(other),
228            DType::F16 => self.assert_eq_elem::<f16>(other),
229            DType::BF16 => self.assert_eq_elem::<bf16>(other),
230            DType::I64 => self.assert_eq_elem::<i64>(other),
231            DType::I32 => self.assert_eq_elem::<i32>(other),
232            DType::I16 => self.assert_eq_elem::<i16>(other),
233            DType::I8 => self.assert_eq_elem::<i8>(other),
234            DType::U64 => self.assert_eq_elem::<u64>(other),
235            DType::U32 => self.assert_eq_elem::<u32>(other),
236            DType::U16 => self.assert_eq_elem::<u16>(other),
237            DType::U8 => self.assert_eq_elem::<u8>(other),
238            DType::Bool(BoolStore::Native) => self.assert_eq_elem::<bool>(other),
239            DType::Bool(BoolStore::U8) => self.assert_eq_elem::<u8>(other),
240            DType::Bool(BoolStore::U32) => self.assert_eq_elem::<u32>(other),
241            DType::QFloat(q) => {
242                // Strict or not, it doesn't make sense to compare quantized data to not quantized data for equality
243                let q_other = if let DType::QFloat(q_other) = other.dtype {
244                    q_other
245                } else {
246                    panic!("Quantized data differs from other not quantized data")
247                };
248
249                // Data equality mostly depends on input quantization type, but we also check level
250                if q.value == q_other.value && q.level == q_other.level {
251                    self.assert_eq_elem::<i8>(other)
252                } else {
253                    panic!("Quantization schemes differ ({q:?} != {q_other:?})")
254                }
255            }
256        }
257    }
258
259    #[track_caller]
260    fn assert_eq_elem<E: Element>(&self, other: &Self) {
261        let mut message = String::new();
262        if self.shape != other.shape {
263            message += format!(
264                "\n  => Shape is different: {:?} != {:?}",
265                self.shape, other.shape
266            )
267            .as_str();
268        }
269
270        let mut num_diff = 0;
271        let max_num_diff = 5;
272        for (i, (a, b)) in self.iter::<E>().zip(other.iter::<E>()).enumerate() {
273            if !a.eq(&b) {
274                // Only print the first 5 different values.
275                if num_diff < max_num_diff {
276                    message += format!("\n  => Position {i}: {a} != {b}").as_str();
277                }
278                num_diff += 1;
279            }
280        }
281
282        if num_diff >= max_num_diff {
283            message += format!("\n{} more errors...", num_diff - max_num_diff).as_str();
284        }
285
286        if !message.is_empty() {
287            panic!("Tensors are not eq:{message}");
288        }
289    }
290
291    /// Asserts the data is approximately equal to another data.
292    ///
293    /// # Arguments
294    ///
295    /// * `other` - The other data.
296    /// * `tolerance` - The tolerance of the comparison.
297    ///
298    /// # Panics
299    ///
300    /// Panics if the data is not approximately equal.
301    #[track_caller]
302    pub fn assert_approx_eq<F: Float + Element>(&self, other: &Self, tolerance: Tolerance<F>) {
303        let mut message = String::new();
304        if self.shape != other.shape {
305            message += format!(
306                "\n  => Shape is different: {:?} != {:?}",
307                self.shape, other.shape
308            )
309            .as_str();
310        }
311
312        let iter = self.iter::<F>().zip(other.iter::<F>());
313
314        let mut num_diff = 0;
315        let max_num_diff = 5;
316
317        for (i, (a, b)) in iter.enumerate() {
318            //if they are both nan, then they are equally nan
319            let both_nan = a.is_nan() && b.is_nan();
320            //this works for both infinities
321            let both_inf =
322                a.is_infinite() && b.is_infinite() && ((a > F::zero()) == (b > F::zero()));
323
324            if both_nan || both_inf {
325                continue;
326            }
327
328            if !tolerance.approx_eq(F::from(a).unwrap(), F::from(b).unwrap()) {
329                // Only print the first 5 different values.
330                if num_diff < max_num_diff {
331                    let diff_abs = ToPrimitive::to_f64(&(a - b).abs()).unwrap();
332                    let max = F::max(a.abs(), b.abs());
333                    let diff_rel = diff_abs / ToPrimitive::to_f64(&max).unwrap();
334
335                    let tol_rel = ToPrimitive::to_f64(&tolerance.relative).unwrap();
336                    let tol_abs = ToPrimitive::to_f64(&tolerance.absolute).unwrap();
337
338                    message += format!(
339                        "\n  => Position {i}: {a} != {b}\n     diff (rel = {diff_rel:+.2e}, abs = {diff_abs:+.2e}), tol (rel = {tol_rel:+.2e}, abs = {tol_abs:+.2e})"
340                    )
341                    .as_str();
342                }
343                num_diff += 1;
344            }
345        }
346
347        if num_diff >= max_num_diff {
348            message += format!("\n{} more errors...", num_diff - 5).as_str();
349        }
350
351        if !message.is_empty() {
352            panic!("Tensors are not approx eq:{message}");
353        }
354    }
355
356    /// Asserts each value is within a given range.
357    ///
358    /// # Arguments
359    ///
360    /// * `range` - The range.
361    ///
362    /// # Panics
363    ///
364    /// If any value is not within the half-open range bounded inclusively below
365    /// and exclusively above (`start..end`).
366    pub fn assert_within_range<E: ElementOrdered>(&self, range: core::ops::Range<E>) {
367        for elem in self.iter::<E>() {
368            if elem.cmp(&range.start).is_lt() || elem.cmp(&range.end).is_ge() {
369                panic!("Element ({elem:?}) is not within range {range:?}");
370            }
371        }
372    }
373
374    /// Asserts each value is within a given inclusive range.
375    ///
376    /// # Arguments
377    ///
378    /// * `range` - The range.
379    ///
380    /// # Panics
381    ///
382    /// If any value is not within the half-open range bounded inclusively (`start..=end`).
383    pub fn assert_within_range_inclusive<E: ElementOrdered>(
384        &self,
385        range: core::ops::RangeInclusive<E>,
386    ) {
387        let start = range.start();
388        let end = range.end();
389
390        for elem in self.iter::<E>() {
391            if elem.cmp(start).is_lt() || elem.cmp(end).is_gt() {
392                panic!("Element ({elem:?}) is not within range {range:?}");
393            }
394        }
395    }
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401
402    #[test]
403    fn should_assert_appox_eq_limit() {
404        let data1 = TensorData::from([[3.0, 5.0, 6.0]]);
405        let data2 = TensorData::from([[3.03, 5.0, 6.0]]);
406
407        data1.assert_approx_eq::<f32>(&data2, Tolerance::absolute(3e-2));
408        data1.assert_approx_eq::<f16>(&data2, Tolerance::absolute(3e-2));
409    }
410
411    #[test]
412    #[should_panic]
413    fn should_assert_approx_eq_above_limit() {
414        let data1 = TensorData::from([[3.0, 5.0, 6.0]]);
415        let data2 = TensorData::from([[3.031, 5.0, 6.0]]);
416
417        data1.assert_approx_eq::<f32>(&data2, Tolerance::absolute(1e-2));
418    }
419
420    #[test]
421    #[should_panic]
422    fn should_assert_approx_eq_check_shape() {
423        let data1 = TensorData::from([[3.0, 5.0, 6.0, 7.0]]);
424        let data2 = TensorData::from([[3.0, 5.0, 6.0]]);
425
426        data1.assert_approx_eq::<f32>(&data2, Tolerance::absolute(1e-2));
427    }
428}
burn_std/data/compare.rs

burn_std/data/
compare.rs
