burn_std/data/
compare.rs

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

burn_std/data/
compare.rs
