tract_data/
datum.rs

1//! `Tensor` is the main data container for tract
2use crate::TVec;
3use crate::dim::TDim;
4use crate::internal::*;
5use crate::tensor::Tensor;
6use half::f16;
7#[cfg(feature = "complex")]
8use num_complex::Complex;
9use std::fmt;
10use std::hash::Hash;
11
12use num_traits::AsPrimitive;
13
14#[derive(Copy, Clone, PartialEq)]
15pub enum QParams {
16    MinMax { min: f32, max: f32 },
17    ZpScale { zero_point: i32, scale: f32 },
18}
19
20impl Eq for QParams {}
21
22impl Ord for QParams {
23    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
24        use QParams::*;
25        match (self, other) {
26            (MinMax { .. }, ZpScale { .. }) => std::cmp::Ordering::Less,
27            (ZpScale { .. }, MinMax { .. }) => std::cmp::Ordering::Greater,
28            (MinMax { min: min1, max: max1 }, MinMax { min: min2, max: max2 }) => {
29                min1.total_cmp(min2).then_with(|| max1.total_cmp(max2))
30            }
31            (
32                Self::ZpScale { zero_point: zp1, scale: s1 },
33                Self::ZpScale { zero_point: zp2, scale: s2 },
34            ) => zp1.cmp(zp2).then_with(|| s1.total_cmp(s2)),
35        }
36    }
37}
38
39impl PartialOrd for QParams {
40    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
41        Some(self.cmp(other))
42    }
43}
44
45impl Default for QParams {
46    fn default() -> Self {
47        QParams::ZpScale { zero_point: 0, scale: 1. }
48    }
49}
50
51#[allow(clippy::derived_hash_with_manual_eq)]
52impl Hash for QParams {
53    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
54        match self {
55            QParams::MinMax { min, max } => {
56                0.hash(state);
57                min.to_bits().hash(state);
58                max.to_bits().hash(state);
59            }
60            QParams::ZpScale { zero_point, scale } => {
61                1.hash(state);
62                zero_point.hash(state);
63                scale.to_bits().hash(state);
64            }
65        }
66    }
67}
68
69impl QParams {
70    pub fn zp_scale(&self) -> (i32, f32) {
71        match self {
72            QParams::MinMax { min, max } => {
73                let scale = (max - min) / 255.;
74                ((-(min + max) / 2. / scale) as i32, scale)
75            }
76            QParams::ZpScale { zero_point, scale } => (*zero_point, *scale),
77        }
78    }
79
80    pub fn q(&self, f: f32) -> i32 {
81        let (zp, scale) = self.zp_scale();
82        (f / scale) as i32 + zp
83    }
84
85    pub fn dq(&self, i: i32) -> f32 {
86        let (zp, scale) = self.zp_scale();
87        (i - zp) as f32 * scale
88    }
89}
90
91impl std::fmt::Debug for QParams {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        let (zp, scale) = self.zp_scale();
94        write!(f, "Z:{zp} S:{scale}")
95    }
96}
97
98#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
99pub enum DatumType {
100    Bool,
101    U8,
102    U16,
103    U32,
104    U64,
105    I8,
106    I16,
107    I32,
108    I64,
109    F16,
110    F32,
111    F64,
112    TDim,
113    Blob,
114    String,
115    QI8(QParams),
116    QU8(QParams),
117    QI32(QParams),
118    #[cfg(feature = "complex")]
119    ComplexI16,
120    #[cfg(feature = "complex")]
121    ComplexI32,
122    #[cfg(feature = "complex")]
123    ComplexI64,
124    #[cfg(feature = "complex")]
125    ComplexF16,
126    #[cfg(feature = "complex")]
127    ComplexF32,
128    #[cfg(feature = "complex")]
129    ComplexF64,
130}
131
132impl DatumType {
133    pub fn super_types(&self) -> TVec<DatumType> {
134        use DatumType::*;
135        if *self == String || *self == TDim || *self == Blob || *self == Bool || self.is_quantized()
136        {
137            return tvec!(*self);
138        }
139        #[cfg(feature = "complex")]
140        if self.is_complex_float() {
141            return [ComplexF16, ComplexF32, ComplexF64]
142                .iter()
143                .filter(|s| s.size_of() >= self.size_of())
144                .copied()
145                .collect();
146        } else if self.is_complex_signed() {
147            return [ComplexI16, ComplexI32, ComplexI64]
148                .iter()
149                .filter(|s| s.size_of() >= self.size_of())
150                .copied()
151                .collect();
152        }
153        if self.is_float() {
154            [F16, F32, F64].iter().filter(|s| s.size_of() >= self.size_of()).copied().collect()
155        } else if self.is_signed() {
156            [I8, I16, I32, I64, TDim]
157                .iter()
158                .filter(|s| s.size_of() >= self.size_of())
159                .copied()
160                .collect()
161        } else {
162            [U8, U16, U32, U64].iter().filter(|s| s.size_of() >= self.size_of()).copied().collect()
163        }
164    }
165
166    pub fn super_type_for(
167        i: impl IntoIterator<Item = impl std::borrow::Borrow<DatumType>>,
168    ) -> Option<DatumType> {
169        let mut iter = i.into_iter();
170        let mut current = {
171            let it = iter.next()?;
172            *it.borrow()
173        };
174        for n in iter {
175            {
176                let it = current.common_super_type(*n.borrow())?;
177                current = it
178            }
179        }
180        Some(current)
181    }
182
183    pub fn common_super_type(&self, rhs: DatumType) -> Option<DatumType> {
184        for mine in self.super_types() {
185            for theirs in rhs.super_types() {
186                if mine == theirs {
187                    return Some(mine);
188                }
189            }
190        }
191        None
192    }
193
194    pub fn is_unsigned(&self) -> bool {
195        matches!(
196            self.unquantized(),
197            DatumType::U8 | DatumType::U16 | DatumType::U32 | DatumType::U64
198        )
199    }
200
201    pub fn is_signed(&self) -> bool {
202        matches!(
203            self.unquantized(),
204            DatumType::I8 | DatumType::I16 | DatumType::I32 | DatumType::I64
205        )
206    }
207
208    pub fn is_float(&self) -> bool {
209        matches!(self, DatumType::F16 | DatumType::F32 | DatumType::F64)
210    }
211
212    pub fn is_number(&self) -> bool {
213        self.is_signed() | self.is_unsigned() | self.is_float() | self.is_quantized()
214    }
215
216    pub fn is_tdim(&self) -> bool {
217        *self == DatumType::TDim
218    }
219
220    #[cfg(feature = "complex")]
221    pub fn is_complex(&self) -> bool {
222        self.is_complex_float() || self.is_complex_signed()
223    }
224
225    #[cfg(feature = "complex")]
226    pub fn is_complex_float(&self) -> bool {
227        matches!(self, DatumType::ComplexF16 | DatumType::ComplexF32 | DatumType::ComplexF64)
228    }
229
230    #[cfg(feature = "complex")]
231    pub fn is_complex_signed(&self) -> bool {
232        matches!(self, DatumType::ComplexI16 | DatumType::ComplexI32 | DatumType::ComplexI64)
233    }
234
235    #[cfg(feature = "complex")]
236    pub fn complexify(&self) -> TractResult<DatumType> {
237        match *self {
238            DatumType::I16 => Ok(DatumType::ComplexI16),
239            DatumType::I32 => Ok(DatumType::ComplexI32),
240            DatumType::I64 => Ok(DatumType::ComplexI64),
241            DatumType::F16 => Ok(DatumType::ComplexF16),
242            DatumType::F32 => Ok(DatumType::ComplexF32),
243            DatumType::F64 => Ok(DatumType::ComplexF64),
244            _ => bail!("No complex datum type formed on {:?}", self),
245        }
246    }
247
248    #[cfg(feature = "complex")]
249    pub fn decomplexify(&self) -> TractResult<DatumType> {
250        match *self {
251            DatumType::ComplexI16 => Ok(DatumType::I16),
252            DatumType::ComplexI32 => Ok(DatumType::I32),
253            DatumType::ComplexI64 => Ok(DatumType::I64),
254            DatumType::ComplexF16 => Ok(DatumType::F16),
255            DatumType::ComplexF32 => Ok(DatumType::F32),
256            DatumType::ComplexF64 => Ok(DatumType::F64),
257            _ => bail!("{:?} is not a complex type", self),
258        }
259    }
260
261    pub fn is_copy(&self) -> bool {
262        #[cfg(feature = "complex")]
263        if self.is_complex() {
264            return true;
265        }
266        *self == DatumType::Bool || self.is_unsigned() || self.is_signed() || self.is_float()
267    }
268
269    pub fn is_quantized(&self) -> bool {
270        self.qparams().is_some()
271    }
272
273    pub fn qparams(&self) -> Option<QParams> {
274        match self {
275            DatumType::QI8(qparams) | DatumType::QU8(qparams) | DatumType::QI32(qparams) => {
276                Some(*qparams)
277            }
278            _ => None,
279        }
280    }
281
282    pub fn with_qparams(&self, qparams: QParams) -> DatumType {
283        match self {
284            DatumType::QI8(_) => DatumType::QI8(qparams),
285            DatumType::QU8(_) => DatumType::QU8(qparams),
286            DatumType::QI32(_) => DatumType::QI32(qparams),
287            _ => *self,
288        }
289    }
290
291    pub fn quantize(&self, qparams: QParams) -> DatumType {
292        match self {
293            DatumType::I8 => DatumType::QI8(qparams),
294            DatumType::U8 => DatumType::QU8(qparams),
295            DatumType::I32 => DatumType::QI32(qparams),
296            DatumType::QI8(_) => DatumType::QI8(qparams),
297            DatumType::QU8(_) => DatumType::QU8(qparams),
298            DatumType::QI32(_) => DatumType::QI32(qparams),
299            _ => panic!("Can't quantize {self:?}"),
300        }
301    }
302
303    #[inline(always)]
304    pub fn zp_scale(&self) -> (i32, f32) {
305        self.qparams().map(|q| q.zp_scale()).unwrap_or((0, 1.))
306    }
307
308    #[inline(always)]
309    pub fn with_zp_scale(&self, zero_point: i32, scale: f32) -> DatumType {
310        self.quantize(QParams::ZpScale { zero_point, scale })
311    }
312
313    pub fn unquantized(&self) -> DatumType {
314        match self {
315            DatumType::QI8(_) => DatumType::I8,
316            DatumType::QU8(_) => DatumType::U8,
317            DatumType::QI32(_) => DatumType::I32,
318            _ => *self,
319        }
320    }
321
322    pub fn integer(signed: bool, size: usize) -> Self {
323        use DatumType::*;
324        match (signed, size) {
325            (false, 8) => U8,
326            (false, 16) => U16,
327            (false, 32) => U32,
328            (false, 64) => U64,
329            (true, 8) => U8,
330            (true, 16) => U16,
331            (true, 32) => U32,
332            (true, 64) => U64,
333            _ => panic!("No integer for signed:{signed} size:{size}"),
334        }
335    }
336
337    pub fn is_integer(&self) -> bool {
338        self.is_signed() || self.is_unsigned()
339    }
340
341    #[inline]
342    pub fn size_of(&self) -> usize {
343        dispatch_datum!(std::mem::size_of(self)())
344    }
345
346    pub fn min_value(&self) -> Tensor {
347        match self {
348            DatumType::QU8(_)
349            | DatumType::U8
350            | DatumType::U16
351            | DatumType::U32
352            | DatumType::U64 => Tensor::zero_dt(*self, &[1]).unwrap(),
353            DatumType::I8 | DatumType::QI8(_) => tensor0(i8::MIN),
354            DatumType::QI32(_) => tensor0(i32::MIN),
355            DatumType::I16 => tensor0(i16::MIN),
356            DatumType::I32 => tensor0(i32::MIN),
357            DatumType::I64 => tensor0(i64::MIN),
358            DatumType::F16 => tensor0(f16::MIN),
359            DatumType::F32 => tensor0(f32::MIN),
360            DatumType::F64 => tensor0(f64::MIN),
361            _ => panic!("No min value for datum type {self:?}"),
362        }
363    }
364    pub fn max_value(&self) -> Tensor {
365        match self {
366            DatumType::U8 | DatumType::QU8(_) => tensor0(u8::MAX),
367            DatumType::U16 => tensor0(u16::MAX),
368            DatumType::U32 => tensor0(u32::MAX),
369            DatumType::U64 => tensor0(u64::MAX),
370            DatumType::I8 | DatumType::QI8(_) => tensor0(i8::MAX),
371            DatumType::I16 => tensor0(i16::MAX),
372            DatumType::I32 => tensor0(i32::MAX),
373            DatumType::I64 => tensor0(i64::MAX),
374            DatumType::QI32(_) => tensor0(i32::MAX),
375            DatumType::F16 => tensor0(f16::MAX),
376            DatumType::F32 => tensor0(f32::MAX),
377            DatumType::F64 => tensor0(f64::MAX),
378            _ => panic!("No max value for datum type {self:?}"),
379        }
380    }
381
382    pub fn is<D: Datum>(&self) -> bool {
383        *self == D::datum_type()
384    }
385}
386
387/// Parses the `Debug` form of a zero-point/scale quantized type, `QU8(Z:128 S:0.01)`.
388fn parse_zp_scale(s: &str, prefix: &str) -> Option<QParams> {
389    let body = s.strip_prefix(prefix)?.strip_prefix("(Z:")?.strip_suffix(")")?;
390    let (zero_point, scale) = body.split_once(" S:")?;
391    Some(QParams::ZpScale { zero_point: zero_point.parse().ok()?, scale: scale.parse().ok()? })
392}
393
394impl std::str::FromStr for DatumType {
395    type Err = TractError;
396
397    fn from_str(s: &str) -> Result<Self, Self::Err> {
398        if let Some(qp) = parse_zp_scale(s, "QU8") {
399            Ok(DatumType::QU8(qp))
400        } else if let Some(qp) = parse_zp_scale(s, "QI8") {
401            Ok(DatumType::QI8(qp))
402        } else if let Some(qp) = parse_zp_scale(s, "QI32") {
403            Ok(DatumType::QI32(qp))
404        } else {
405            match s {
406                "I8" | "i8" => Ok(DatumType::I8),
407                "I16" | "i16" => Ok(DatumType::I16),
408                "I32" | "i32" => Ok(DatumType::I32),
409                "I64" | "i64" => Ok(DatumType::I64),
410                "U8" | "u8" => Ok(DatumType::U8),
411                "U16" | "u16" => Ok(DatumType::U16),
412                "U32" | "u32" => Ok(DatumType::U32),
413                "U64" | "u64" => Ok(DatumType::U64),
414                "F16" | "f16" => Ok(DatumType::F16),
415                "F32" | "f32" => Ok(DatumType::F32),
416                "F64" | "f64" => Ok(DatumType::F64),
417                "Bool" | "bool" => Ok(DatumType::Bool),
418                "Blob" | "blob" => Ok(DatumType::Blob),
419                "String" | "string" => Ok(DatumType::String),
420                "TDim" | "tdim" => Ok(DatumType::TDim),
421                #[cfg(feature = "complex")]
422                "ComplexI16" | "complexi16" => Ok(DatumType::ComplexI16),
423                #[cfg(feature = "complex")]
424                "ComplexI32" | "complexi32" => Ok(DatumType::ComplexI32),
425                #[cfg(feature = "complex")]
426                "ComplexI64" | "complexi64" => Ok(DatumType::ComplexI64),
427                #[cfg(feature = "complex")]
428                "ComplexF16" | "complexf16" => Ok(DatumType::ComplexF16),
429                #[cfg(feature = "complex")]
430                "ComplexF32" | "complexf32" => Ok(DatumType::ComplexF32),
431                #[cfg(feature = "complex")]
432                "ComplexF64" | "complexf64" => Ok(DatumType::ComplexF64),
433                _ => bail!("Unknown type {}", s),
434            }
435        }
436    }
437}
438
439const TOINT: f32 = 1.0f32 / f32::EPSILON;
440
441pub fn round_ties_to_even(x: f32) -> f32 {
442    let u = x.to_bits();
443    let e = (u >> 23) & 0xff;
444    if e >= 0x7f + 23 {
445        return x;
446    }
447    let s = u >> 31;
448    let y = if s == 1 { x - TOINT + TOINT } else { x + TOINT - TOINT };
449    if y == 0.0 { if s == 1 { -0f32 } else { 0f32 } } else { y }
450}
451
452#[inline]
453pub fn scale_by<T: Datum + AsPrimitive<f32>>(b: T, a: f32) -> T
454where
455    f32: AsPrimitive<T>,
456{
457    let b = b.as_();
458    (round_ties_to_even(b.abs() * a) * b.signum()).as_()
459}
460
461pub trait ClampCast: PartialOrd + Copy + 'static {
462    #[inline(always)]
463    fn clamp_cast<O>(self) -> O
464    where
465        Self: AsPrimitive<O> + Datum,
466        O: AsPrimitive<Self> + num_traits::Bounded + Datum,
467    {
468        // this fails if we're upcasting, in which case clamping is useless
469        if O::min_value().as_() < O::max_value().as_() {
470            num_traits::clamp(self, O::min_value().as_(), O::max_value().as_()).as_()
471        } else {
472            self.as_()
473        }
474    }
475}
476impl<T: PartialOrd + Copy + 'static> ClampCast for T {}
477
478pub trait Datum:
479    Clone + Send + Sync + fmt::Debug + fmt::Display + Default + 'static + PartialEq
480{
481    fn name() -> &'static str;
482    fn datum_type() -> DatumType;
483    fn is<D: Datum>() -> bool;
484}
485
486macro_rules! datum {
487    ($t:ty, $v:ident) => {
488        impl From<$t> for Tensor {
489            fn from(it: $t) -> Tensor {
490                tensor0(it)
491            }
492        }
493
494        impl Datum for $t {
495            fn name() -> &'static str {
496                stringify!($t)
497            }
498
499            fn datum_type() -> DatumType {
500                DatumType::$v
501            }
502
503            fn is<D: Datum>() -> bool {
504                Self::datum_type() == D::datum_type()
505            }
506        }
507    };
508}
509
510datum!(bool, Bool);
511datum!(f16, F16);
512datum!(f32, F32);
513datum!(f64, F64);
514datum!(i8, I8);
515datum!(i16, I16);
516datum!(i32, I32);
517datum!(i64, I64);
518datum!(u8, U8);
519datum!(u16, U16);
520datum!(u32, U32);
521datum!(u64, U64);
522datum!(TDim, TDim);
523datum!(String, String);
524datum!(crate::blob::Blob, Blob);
525#[cfg(feature = "complex")]
526datum!(Complex<i16>, ComplexI16);
527#[cfg(feature = "complex")]
528datum!(Complex<i32>, ComplexI32);
529#[cfg(feature = "complex")]
530datum!(Complex<i64>, ComplexI64);
531#[cfg(feature = "complex")]
532datum!(Complex<f16>, ComplexF16);
533#[cfg(feature = "complex")]
534datum!(Complex<f32>, ComplexF32);
535#[cfg(feature = "complex")]
536datum!(Complex<f64>, ComplexF64);
537
538#[cfg(test)]
539mod tests {
540    use crate::internal::*;
541    use ndarray::arr1;
542
543    #[test]
544    fn test_array_to_tensor_to_array() {
545        let array = arr1(&[12i32, 42]);
546        let tensor = Tensor::from(array.clone());
547        let view = tensor.to_plain_array_view::<i32>().unwrap();
548        assert_eq!(array, view.into_dimensionality().unwrap());
549    }
550
551    #[test]
552    fn test_cast_dim_to_dim() {
553        let t_dim: Tensor = tensor1(&[12isize.to_dim(), 42isize.to_dim()]);
554        let t_i32 = t_dim.cast_to::<i32>().unwrap();
555        let t_dim_2 = t_i32.cast_to::<TDim>().unwrap().into_owned();
556        assert_eq!(t_dim, t_dim_2);
557    }
558
559    #[test]
560    fn test_cast_i32_to_dim() {
561        let t_i32: Tensor = tensor1(&[0i32, 12]);
562        t_i32.cast_to::<TDim>().unwrap();
563    }
564
565    #[test]
566    fn test_cast_i64_to_bool() {
567        let t_i64: Tensor = tensor1(&[0i64]);
568        t_i64.cast_to::<bool>().unwrap();
569    }
570
571    #[test]
572    fn test_parse_qu8() {
573        assert_eq!(
574            "QU8(Z:128 S:0.01)".parse::<DatumType>().unwrap(),
575            DatumType::QU8(QParams::ZpScale { zero_point: 128, scale: 0.01 })
576        );
577    }
578
579    #[test]
580    fn test_parse_quantized_round_trip() {
581        for dt in [
582            DatumType::QI8(QParams::ZpScale { zero_point: -3, scale: 1. }),
583            DatumType::QI32(QParams::ZpScale { zero_point: 0, scale: 2.5e-3 }),
584        ] {
585            assert_eq!(format!("{dt:?}").parse::<DatumType>().unwrap(), dt);
586        }
587    }
588
589    #[test]
590    fn test_parse_malformed_quantized() {
591        for spec in ["QU8(Z:128 S:)", "QU8(Z:128)", "QU8(Z:x S:0.01)", "QU8(Z:128 S:0.01"] {
592            assert!(spec.parse::<DatumType>().is_err());
593        }
594    }
595}
tract_data/datum.rs

tract_data/
datum.rs
