Skip to main content

tract_data/
tensor.rs

1//! `Tensor`, tract main data object of interest.
2use crate::TVec;
3use crate::blob::Blob;
4use crate::datum::{ClampCast, Datum, DatumType, QParams, round_ties_to_even, scale_by};
5use crate::dim::TDim;
6use crate::internal::*;
7use half::f16;
8use itertools::{Itertools, izip};
9use ndarray::prelude::*;
10#[cfg(feature = "complex")]
11use num_complex::Complex;
12use num_traits::Float;
13use std::borrow::Cow;
14use std::fmt;
15use std::hash::Hash;
16use std::ops::Range;
17use std::sync::Arc;
18
19pub mod litteral;
20pub mod plain_view;
21pub mod storage;
22pub mod view;
23
24pub use plain_view::{PlainView, PlainViewMut};
25use storage::{PlainStorage, StorageKind, TensorStorage};
26
27#[derive(Copy, Clone, Default, Debug)]
28pub enum Approximation {
29    Exact,
30    #[default]
31    Close,
32    Approximate,
33    VeryApproximate,
34    SuperApproximate,
35    UltraApproximate,
36    Custom(f32, f32, f32),
37    /// Compare by integer ULP distance in the reference tensor's own float type,
38    /// accepting a distance up to the given bound.
39    ///
40    /// Unlike the tolerance-based variants this does not go through an f32 cast,
41    /// so an f16 comparison stays an f16 comparison. Use it to assert that two
42    /// implementations of a kernel agree to within a known number of rounding
43    /// steps.
44    Ulp(u64),
45}
46
47impl PartialEq for Approximation {
48    fn eq(&self, other: &Self) -> bool {
49        use Approximation::*;
50        match (self, other) {
51            (Custom(aa, ar, ao), Custom(ba, br, bo)) => aa == ba && ar == br && bo == ao,
52            (Ulp(a), Ulp(b)) => a == b,
53            _ => std::mem::discriminant(self) == std::mem::discriminant(other),
54        }
55    }
56}
57
58impl Eq for Approximation {}
59
60impl From<bool> for Approximation {
61    fn from(b: bool) -> Self {
62        if b { Self::Approximate } else { Self::Exact }
63    }
64}
65
66impl Approximation {
67    fn atol_rtol_outliers(&self, dt: &DatumType) -> (f64, f64, f64) {
68        use Approximation::*;
69        match (self, dt) {
70            (Exact, _) => (0.0, 0.0, 0.0),
71            (Close, DatumType::F16) => (1e-3, 1e-3, 0.0),
72            (Approximate, DatumType::F16) => (1e-3, 5e-3, 0.0),
73            (Approximate, qp) if qp.is_quantized() => (qp.zp_scale().1 as f64, 0., 0.0),
74            (Close, _) => (1e-7, 1e-7, 0.0),
75            (Approximate, _) => (1e-4, 5e-4, 0.0),
76            (VeryApproximate, _) => (5e-2, 1e-2, 0.0),
77            (SuperApproximate, _) => (0.1, 0.05, 0.0001),
78            (UltraApproximate, _) => (0.2, 0.1, 0.0005),
79            (Custom(atol, rtol, out), _) => (*atol as _, *rtol as _, *out as _),
80            // Handled by a dedicated path in `Tensor::close_enough`; these values
81            // are never consulted.
82            (Ulp(_), _) => (0.0, 0.0, 0.0),
83        }
84    }
85}
86
87/// Tensor is a concrete tensor in tract.
88pub struct Tensor {
89    dt: DatumType,
90    shape: TVec<usize>,
91    strides: TVec<isize>,
92    len: usize,
93    storage: StorageKind,
94}
95
96unsafe impl Send for Tensor {}
97unsafe impl Sync for Tensor {}
98
99impl Hash for Tensor {
100    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
101        use DatumType::*;
102        self.dt.hash(state);
103        self.shape.hash(state);
104        if let Some(plain) = self.storage.as_plain() {
105            plain.layout().align().hash(state);
106            unsafe {
107                match self.dt {
108                    Bool => self.as_slice_unchecked::<bool>().hash(state),
109                    I8 => self.as_slice_unchecked::<i8>().hash(state),
110                    I16 => self.as_slice_unchecked::<i16>().hash(state),
111                    I32 => self.as_slice_unchecked::<i32>().hash(state),
112                    I64 => self.as_slice_unchecked::<i64>().hash(state),
113                    U8 => self.as_slice_unchecked::<u8>().hash(state),
114                    U16 => self.as_slice_unchecked::<u16>().hash(state),
115                    U32 => self.as_slice_unchecked::<u32>().hash(state),
116                    U64 => self.as_slice_unchecked::<u64>().hash(state),
117                    F16 => self.as_slice_unchecked::<i16>().hash(state),
118                    F32 => self.as_slice_unchecked::<i32>().hash(state),
119                    F64 => self.as_slice_unchecked::<i64>().hash(state),
120                    TDim => self.as_slice_unchecked::<crate::dim::TDim>().hash(state),
121                    String => self.as_slice_unchecked::<std::string::String>().hash(state),
122                    Blob => self.as_slice_unchecked::<crate::blob::Blob>().hash(state),
123                    QI8(_) => self.as_slice_unchecked::<i8>().hash(state),
124                    QU8(_) => self.as_slice_unchecked::<u8>().hash(state),
125                    QI32(_) => self.as_slice_unchecked::<i32>().hash(state),
126                    #[cfg(feature = "complex")]
127                    ComplexI16 => self.as_slice_unchecked::<Complex<i16>>().hash(state),
128                    #[cfg(feature = "complex")]
129                    ComplexI32 => self.as_slice_unchecked::<Complex<i32>>().hash(state),
130                    #[cfg(feature = "complex")]
131                    ComplexI64 => self.as_slice_unchecked::<Complex<i64>>().hash(state),
132                    #[cfg(feature = "complex")]
133                    ComplexF16 => self.as_slice_unchecked::<Complex<i16>>().hash(state),
134                    #[cfg(feature = "complex")]
135                    ComplexF32 => self.as_slice_unchecked::<Complex<i32>>().hash(state),
136                    #[cfg(feature = "complex")]
137                    ComplexF64 => self.as_slice_unchecked::<Complex<i64>>().hash(state),
138                }
139            }
140        } else {
141            self.storage.dyn_hash(state);
142        }
143    }
144}
145
146impl Clone for Tensor {
147    fn clone(&self) -> Tensor {
148        self.deep_clone()
149    }
150}
151
152impl Default for Tensor {
153    fn default() -> Tensor {
154        litteral::tensor0(0f32)
155    }
156}
157
158impl Drop for Tensor {
159    fn drop(&mut self) {
160        if self.is_plain() {
161            macro_rules! drop_in_place {
162                ($t: ty) => {
163                    if self.dt == <$t>::datum_type() {
164                        unsafe {
165                            let slice = self.as_slice_mut_unchecked::<$t>();
166                            std::ptr::drop_in_place(slice as *mut [$t]);
167                        }
168                    }
169                };
170            }
171            drop_in_place!(Blob);
172            drop_in_place!(String);
173            drop_in_place!(TDim);
174        }
175        // StorageKind::Exotic drops via Box<dyn TensorStorage> automatically
176    }
177}
178
179#[allow(unreachable_code)]
180pub fn vector_size() -> usize {
181    #[cfg(target_arch = "x86_64")]
182    {
183        return if is_x86_feature_detected!("avx512f") { 512 / 8 } else { 256 / 8 };
184    }
185    128 / 8
186}
187
188/// Copy `outer` blocks of `block` bytes from a contiguous source into a destination
189/// strided by `out_stride`, as `T`-sized items. Used for blocks too small for a
190/// `copy_nonoverlapping` call per block to pay for itself.
191///
192/// # Safety
193/// `block` and `out_stride` must be multiples of `size_of::<T>()`, both pointers must
194/// be `T`-aligned, and the two ranges must not overlap.
195#[inline]
196unsafe fn copy_blocks<T: Copy>(
197    src: *const u8,
198    dst: *mut u8,
199    outer: usize,
200    block: usize,
201    out_stride: usize,
202) {
203    unsafe {
204        let n = block / std::mem::size_of::<T>();
205        for o in 0..outer {
206            let s = src.add(o * block) as *const T;
207            let d = dst.add(o * out_stride) as *mut T;
208            for i in 0..n {
209                *d.add(i) = *s.add(i);
210            }
211        }
212    }
213}
214
215impl Tensor {
216    #[inline]
217    fn plain_storage(&self) -> &PlainStorage {
218        self.storage.as_plain().expect("Non-plain storage")
219    }
220
221    #[inline]
222    fn plain_storage_mut(&mut self) -> &mut PlainStorage {
223        self.storage.as_plain_mut().expect("Non-plain storage")
224    }
225
226    pub fn storage_as<T: TensorStorage>(&self) -> Option<&T> {
227        self.storage.as_storage().downcast_ref::<T>()
228    }
229
230    pub fn try_storage_as<T: TensorStorage>(&self) -> TractResult<&T> {
231        self.storage_as::<T>().context("Unexpected tensor storage type")
232    }
233
234    pub fn from_storage(
235        dt: DatumType,
236        shape: &[usize],
237        storage: impl TensorStorage + 'static,
238    ) -> Tensor {
239        let len = shape.iter().product::<usize>();
240        let strides = Self::natural_strides(shape);
241        Tensor {
242            dt,
243            shape: shape.into(),
244            strides,
245            len,
246            storage: StorageKind::Exotic(Box::new(storage)),
247        }
248    }
249
250    /// Returns an immutable [`PlainView`] if this tensor has plain storage.
251    #[inline]
252    pub fn as_plain(&self) -> Option<PlainView<'_>> {
253        let storage = self.storage.as_plain()?;
254        Some(PlainView::new(self, storage))
255    }
256
257    /// Returns an immutable [`PlainView`], or an error if storage is not plain.
258    #[inline]
259    pub fn try_as_plain(&self) -> TractResult<PlainView<'_>> {
260        self.as_plain().context("Tensor storage is not plain")
261    }
262
263    /// Returns `true` if this tensor uses plain (contiguous) storage.
264    #[inline]
265    pub fn is_plain(&self) -> bool {
266        self.storage.as_plain().is_some()
267    }
268
269    /// Returns `true` if this tensor uses exotic (non-plain) storage.
270    #[inline]
271    pub fn is_exotic(&self) -> bool {
272        !self.is_plain()
273    }
274
275    /// Build the `ExoticFact` matching this tensor's storage, or `None` for plain tensors.
276    pub fn exotic_fact(&self) -> TractResult<Option<Box<dyn crate::exotic::ExoticFact>>> {
277        self.storage.as_storage().exotic_fact(&self.shape)
278    }
279
280    /// Returns a mutable [`PlainViewMut`] if this tensor has plain storage.
281    #[inline]
282    pub fn as_plain_mut(&mut self) -> Option<PlainViewMut<'_>> {
283        let storage = self.storage.as_plain_mut()?;
284        Some(PlainViewMut::new(self.dt, &self.shape, &self.strides, self.len, storage))
285    }
286
287    /// Returns a mutable [`PlainViewMut`], or an error if storage is not plain.
288    #[inline]
289    pub fn try_as_plain_mut(&mut self) -> TractResult<PlainViewMut<'_>> {
290        self.as_plain_mut().context("Tensor storage is not plain")
291    }
292
293    /// Create an uninitialized tensor (dt as type paramater).
294    #[inline]
295    pub unsafe fn uninitialized<T: Datum>(shape: &[usize]) -> TractResult<Tensor> {
296        unsafe { Self::uninitialized_dt(T::datum_type(), shape) }
297    }
298
299    /// Create an uninitialized tensor (dt as regular parameter).
300    #[inline]
301    pub unsafe fn uninitialized_dt(dt: DatumType, shape: &[usize]) -> TractResult<Tensor> {
302        unsafe { Self::uninitialized_aligned_dt(dt, shape, vector_size()) }
303    }
304
305    /// Create an uninitialized tensor with a given alignment (in bytes).
306    #[inline]
307    pub unsafe fn uninitialized_aligned<T: Datum>(
308        shape: &[usize],
309        alignment: usize,
310    ) -> TractResult<Tensor> {
311        unsafe { Self::uninitialized_aligned_dt(T::datum_type(), shape, alignment) }
312    }
313
314    /// Create an uninitialized tensor with a given alignment (in bytes).
315    pub unsafe fn uninitialized_aligned_dt(
316        dt: DatumType,
317        shape: &[usize],
318        alignment: usize,
319    ) -> TractResult<Tensor> {
320        // `shape` and `dt` come from the model file. Computing the byte count
321        // with a plain product used to wrap around on overflow, silently
322        // allocating a buffer much smaller than the tensor claims (or asking
323        // the allocator for an absurd one). Check instead.
324        let bytes = shape
325            .iter()
326            .try_fold(dt.size_of(), |acc, &d| acc.checked_mul(d))
327            .filter(|&b| b <= isize::MAX as usize)
328            .ok_or_else(|| format_err!("tensor shape {shape:?} of {dt:?} is too large"))?;
329        let storage = StorageKind::Plain(PlainStorage::from(unsafe {
330            Blob::new_for_size_and_align(bytes, alignment)
331        }));
332        let mut tensor = Tensor { strides: tvec!(), dt, shape: shape.into(), storage, len: 0 };
333        if tensor.shape.len() == 0 {
334            tensor.len = 1;
335        } else {
336            tensor.update_strides_and_len();
337        }
338        if !tensor.storage.is_empty() {
339            unsafe fn write_defaults<T: Datum + Default>(tensor: &mut Tensor) {
340                unsafe {
341                    let len = tensor.len;
342                    let dst = tensor.as_slice_mut_unchecked::<T>().as_mut_ptr();
343                    for i in 0..len {
344                        std::ptr::write(dst.add(i), T::default());
345                    }
346                }
347            }
348            if dt == String::datum_type() {
349                unsafe { write_defaults::<String>(&mut tensor) }
350            } else if dt == Blob::datum_type() {
351                unsafe { write_defaults::<Blob>(&mut tensor) }
352            } else if dt == TDim::datum_type() {
353                unsafe { write_defaults::<TDim>(&mut tensor) }
354            } else if cfg!(debug_assertions) {
355                assert!(dt.is_copy());
356                if dt == DatumType::F32 {
357                    tensor.fill_t(f32::NAN).unwrap();
358                } else {
359                    // safe, non copy types have been dealt with
360                    tensor.as_bytes_mut().iter_mut().for_each(|x| *x = (-1i8) as u8);
361                }
362            }
363        }
364        Ok(tensor)
365    }
366
367    pub fn stack_tensors(
368        axis: usize,
369        tensors: &[impl std::borrow::Borrow<Tensor>],
370    ) -> TractResult<Tensor> {
371        ensure!(tensors.len() > 0);
372        let rank = tensors[0].borrow().rank();
373        ensure!(axis < rank);
374        ensure!(tensors.iter().all(|t| t.borrow().rank() == rank));
375        let dt = tensors[0].borrow().datum_type();
376        ensure!(tensors.iter().all(|t| t.borrow().datum_type() == dt));
377        let mut shape: TVec<usize> = tensors[0].borrow().shape().into();
378        for ax in 0..rank {
379            if ax != axis {
380                ensure!(tensors.iter().all(|t| t.borrow().shape()[ax] == shape[ax]));
381            }
382        }
383        shape[axis] = tensors.iter().map(|v| v.borrow().shape()[axis]).sum();
384        unsafe {
385            let mut result = Tensor::uninitialized_dt(dt, &shape)?;
386            // Every input keeps the same trailing block, so one outer stride walks
387            // them alongside the result and each contribution stays contiguous.
388            let outer: usize = shape[..axis].iter().product();
389            let out_stride = shape[axis..].iter().product::<usize>() * dt.size_of();
390            // Each contribution is `outer` blocks of `block` bytes, strided by
391            // `out_stride` in the result. At one f32 per block -- DTLN's
392            // [1, 2, 128, 2] axis-3 concat, FastEnhancer's [1, 256, 1, 2] -- a
393            // copy_nonoverlapping per block is `outer` calls to move four bytes
394            // each, and the generic strided assign below is no better. Copy those
395            // inline, typed, instead: no call, and the loop is a plain strided
396            // store LLVM can widen.
397            const SMALL_BLOCK_BYTES: usize = 64;
398            if dt.is_copy()
399                && outer > 0
400                && tensors.iter().all(|t| t.borrow().storage.as_plain().is_some())
401            {
402                let out = result.plain_storage_mut().as_mut_ptr();
403                let mut offset = 0isize;
404                for v in tensors {
405                    let v = v.borrow();
406                    let block = v.storage.byte_len() / outer;
407                    let src = v.plain_storage().as_ptr();
408                    let dst = out.offset(offset);
409                    if outer == 1 {
410                        std::ptr::copy_nonoverlapping(src, dst, block);
411                    } else if block >= SMALL_BLOCK_BYTES {
412                        for o in 0..outer {
413                            std::ptr::copy_nonoverlapping(
414                                src.add(o * block),
415                                dst.add(o * out_stride),
416                                block,
417                            );
418                        }
419                    } else {
420                        // `block` and both pointers are multiples of the datum size,
421                        // so the typed copy stays aligned.
422                        match dt.size_of() {
423                            1 => copy_blocks::<u8>(src, dst, outer, block, out_stride),
424                            2 => copy_blocks::<u16>(src, dst, outer, block, out_stride),
425                            4 => copy_blocks::<u32>(src, dst, outer, block, out_stride),
426                            8 => copy_blocks::<u64>(src, dst, outer, block, out_stride),
427                            16 => copy_blocks::<u128>(src, dst, outer, block, out_stride),
428                            _ => {
429                                for o in 0..outer {
430                                    std::ptr::copy_nonoverlapping(
431                                        src.add(o * block),
432                                        dst.add(o * out_stride),
433                                        block,
434                                    );
435                                }
436                            }
437                        }
438                    }
439                    offset += block as isize;
440                }
441            } else {
442                let mut offset = 0;
443                for t in tensors {
444                    let t = t.borrow();
445                    let len = t.shape()[axis];
446                    result.assign_slice_from_resolved(
447                        &[],
448                        offset..offset + len,
449                        t,
450                        &[],
451                        0..len,
452                        axis,
453                    );
454                    offset += len;
455                }
456            }
457
458            Ok(result)
459        }
460    }
461
462    pub fn clear<T: Datum + num_traits::Zero + Clone>(&mut self) -> TractResult<()> {
463        self.fill_t(T::zero())
464    }
465
466    pub fn zero<T: Datum + num_traits::Zero>(shape: &[usize]) -> TractResult<Tensor> {
467        unsafe {
468            let mut t = Tensor::uninitialized::<T>(shape)?;
469            t.clear::<T>()?;
470            Ok(t)
471        }
472    }
473
474    pub fn zero_scalar<T: Datum + num_traits::Zero>() -> TractResult<Tensor> {
475        Tensor::zero::<T>(&[])
476    }
477
478    pub fn zero_scalar_dt(dt: DatumType) -> TractResult<Tensor> {
479        Tensor::zero_dt(dt, &[])
480    }
481
482    pub fn zero_dt(dt: DatumType, shape: &[usize]) -> TractResult<Tensor> {
483        Tensor::zero_aligned_dt(dt, shape, vector_size())
484    }
485
486    pub fn fill_t<T: Datum + Clone>(&mut self, value: T) -> TractResult<()> {
487        self.try_as_plain_mut()?
488            .as_slice_mut::<T>()?
489            .iter_mut()
490            .for_each(|item| *item = value.clone());
491        Ok(())
492    }
493
494    pub fn zero_aligned_dt(
495        dt: DatumType,
496        shape: &[usize],
497        alignment: usize,
498    ) -> TractResult<Tensor> {
499        if shape.iter().product::<usize>() == 0 {
500            unsafe { return Tensor::uninitialized_dt(dt, shape) };
501        }
502        if dt.is_quantized() {
503            unsafe {
504                let mut t = Tensor::uninitialized_dt(dt, shape)?;
505                let zp = dt.zp_scale().0;
506                match dt.unquantized() {
507                    DatumType::I8 => t
508                        .try_as_plain_mut()?
509                        .as_slice_mut::<i8>()?
510                        .iter_mut()
511                        .for_each(|item| *item = zp as _),
512                    DatumType::U8 => t
513                        .try_as_plain_mut()?
514                        .as_slice_mut::<u8>()?
515                        .iter_mut()
516                        .for_each(|item| *item = zp as _),
517                    DatumType::I32 => t
518                        .try_as_plain_mut()?
519                        .as_slice_mut::<i32>()?
520                        .iter_mut()
521                        .for_each(|item| *item = zp as _),
522                    _ => unreachable!(),
523                }
524                Ok(t)
525            }
526        } else if dt == DatumType::Bool {
527            let mut t = unsafe { Tensor::uninitialized_dt(dt, shape)? };
528            t.fill_t::<bool>(false)?;
529            Ok(t)
530        } else {
531            dispatch_zerolike!(Self::zero_aligned(dt)(shape, alignment))
532        }
533    }
534
535    pub fn zero_aligned<T: Datum + num_traits::Zero>(
536        shape: &[usize],
537        alignment: usize,
538    ) -> TractResult<Tensor> {
539        unsafe {
540            let mut tensor = Self::uninitialized_aligned::<T>(shape, alignment)?;
541            tensor.clear::<T>()?;
542            Ok(tensor)
543        }
544    }
545
546    /// Create a tensor with a given shape and a slice of elements.
547    /// The data is copied and aligned to size of T.
548    pub fn from_shape<T: Datum + Copy>(shape: &[usize], data: &[T]) -> TractResult<Tensor> {
549        Self::from_shape_align(shape, data, vector_size())
550    }
551
552    /// Create a tensor with a given shape and a slice of elements.
553    /// The data is copied and aligned to given alignment.
554    pub fn from_shape_align<T: Datum + Copy>(
555        shape: &[usize],
556        data: &[T],
557        align: usize,
558    ) -> TractResult<Tensor> {
559        ensure!(
560            data.len() == shape.iter().product::<usize>(),
561            "Shape product must be equal to data length"
562        );
563        unsafe {
564            let bytes = std::slice::from_raw_parts(
565                data.as_ptr() as *const u8,
566                data.len() * T::datum_type().size_of(),
567            );
568            let dt = T::datum_type();
569            Self::from_raw_dt_align(dt, shape, bytes, align)
570        }
571    }
572
573    /// Create a tensor from raw data.
574    ///
575    /// It copies the data, aligning it to the size of T.
576    pub unsafe fn from_raw<T: Datum>(shape: &[usize], content: &[u8]) -> TractResult<Tensor> {
577        unsafe { Tensor::from_raw_dt(T::datum_type(), shape, content) }
578    }
579
580    pub unsafe fn from_raw_aligned<T: Datum>(
581        shape: &[usize],
582        content: &[u8],
583        align: usize,
584    ) -> TractResult<Tensor> {
585        unsafe { Tensor::from_raw_dt_align(T::datum_type(), shape, content, align) }
586    }
587
588    pub unsafe fn from_raw_dt(
589        dt: DatumType,
590        shape: &[usize],
591        content: &[u8],
592    ) -> TractResult<Tensor> {
593        unsafe { Self::from_raw_dt_align(dt, shape, content, vector_size()) }
594    }
595
596    pub unsafe fn from_raw_dt_align(
597        dt: DatumType,
598        shape: &[usize],
599        content: &[u8],
600        align: usize,
601    ) -> TractResult<Tensor> {
602        // Check the declared shape against the payload *before* allocating.
603        // The shape is attacker-controlled, so allocating first lets a
604        // malformed model request an allocation of any size it likes, even
605        // when the payload that follows is a few bytes long.
606        let len = shape
607            .iter()
608            .try_fold(1usize, |acc, &d| acc.checked_mul(d))
609            .ok_or_else(|| format_err!("tensor shape {shape:?} overflows"))?;
610        let expected = len
611            .checked_mul(dt.size_of())
612            .ok_or_else(|| format_err!("tensor shape {shape:?} of {dt:?} is too large"))?;
613        ensure!(
614            content.len() == expected,
615            "Raw tensor data length ({}) does not match shape {:?} of {:?} ({} bytes)",
616            content.len(),
617            shape,
618            dt,
619            expected
620        );
621        let mut tensor = unsafe { Tensor::uninitialized_aligned_dt(dt, shape, align) }?;
622        tensor.as_bytes_mut().copy_from_slice(content);
623        Ok(tensor)
624    }
625
626    pub unsafe fn from_slice_align<T: Datum>(content: &[T], align: usize) -> TractResult<Tensor> {
627        let bytes = if content.len() == 0 {
628            &[]
629        } else {
630            unsafe {
631                std::slice::from_raw_parts(
632                    content.as_ptr() as *const u8,
633                    content.len() * T::datum_type().size_of(),
634                )
635            }
636        };
637        unsafe { Self::from_raw_dt_align(T::datum_type(), &[content.len()], bytes, align) }
638    }
639
640    /// Get the number of dimensions (or axes) of the tensor.
641    #[inline]
642    pub fn rank(&self) -> usize {
643        self.shape.len()
644    }
645
646    /// Get the shape of the tensor.
647    #[inline]
648    pub fn shape(&self) -> &[usize] {
649        &self.shape
650    }
651
652    /// Get the number of values in the tensor.
653    #[inline]
654    #[allow(clippy::len_without_is_empty)]
655    pub fn len(&self) -> usize {
656        self.len
657    }
658
659    /// Get the number of valeus in the tensor.
660    #[inline]
661    #[allow(clippy::len_without_is_empty)]
662    pub fn volume(&self) -> usize {
663        self.len
664    }
665
666    /// Get the shape of the tensor.
667    #[inline]
668    pub fn strides(&self) -> &[isize] {
669        &self.strides
670    }
671
672    fn update_strides_and_len(&mut self) {
673        self.strides.clear();
674        if self.shape.len() == 0 {
675            self.len = 1;
676            return;
677        }
678        compute_natural_stride_to(&mut self.strides, &self.shape);
679        self.len = unsafe { *self.strides.get_unchecked(0) as usize * self.shape.get_unchecked(0) };
680    }
681
682    /// Force the tensor shape, no consistency check.
683    pub unsafe fn set_shape_unchecked(&mut self, shape: &[usize]) {
684        if shape != &*self.shape {
685            self.shape.clear();
686            self.shape.extend_from_slice(shape);
687            self.update_strides_and_len();
688        }
689    }
690
691    /// Force the tensor shape and strides, no consistency check.
692    pub unsafe fn set_geometry_unchecked(&mut self, shape: &[usize], strides: &[isize]) {
693        self.shape.clear();
694        self.shape.extend_from_slice(shape);
695        self.strides.clear();
696        self.strides.extend_from_slice(strides);
697    }
698
699    /// Force the tensor shape.
700    pub fn set_shape(&mut self, shape: &[usize]) -> TractResult<()> {
701        if self.len() != shape.iter().product::<usize>() {
702            bail!("Invalid reshape {:?} to {:?}", self.shape, shape);
703        }
704        unsafe { self.set_shape_unchecked(shape) }
705        Ok(())
706    }
707
708    pub fn permute_axes(self, axes: &[usize]) -> TractResult<Tensor> {
709        ensure!(axes.iter().duplicates().next().is_none());
710        ensure!(axes.iter().all(|a| *a < self.rank()));
711        unsafe {
712            #[inline]
713            unsafe fn permute<T: Datum>(axes: &[usize], input: Tensor) -> Tensor {
714                unsafe { input.into_array_unchecked::<T>().permuted_axes(axes).into_tensor() }
715            }
716            let dt = self.datum_type();
717            let mut t = dispatch_datum_by_size!(permute(self.datum_type())(axes, self));
718            t.set_datum_type(dt);
719            Ok(t)
720        }
721    }
722
723    pub fn move_axis(self, from: usize, to: usize) -> TractResult<Tensor> {
724        let mut permutation: Vec<usize> = (0..self.rank()).collect();
725        permutation.remove(from);
726        permutation.insert(to, from);
727        self.permute_axes(&permutation)
728    }
729
730    pub fn collapse_axis_with_next(mut self, axis: usize) -> Tensor {
731        let removed = self.shape.remove(axis + 1);
732        self.shape[axis] *= removed;
733        self.update_strides_and_len();
734        self
735    }
736
737    pub fn split_axis(mut self, axis: usize, outer_dim: usize) -> TractResult<Tensor> {
738        if !self.shape[axis].is_multiple_of(outer_dim) {
739            bail!(
740                "Invalid axis split, shape is {:?}, axis split at {}, outer {}",
741                self.shape,
742                axis,
743                outer_dim
744            );
745        }
746        self.shape.insert(axis + 1, self.shape[axis] / outer_dim);
747        self.shape[axis] = outer_dim;
748        self.update_strides_and_len();
749        Ok(self)
750    }
751
752    /// Reshape the tensor to `shape`.
753    pub fn into_shape(mut self, shape: &[usize]) -> TractResult<Tensor> {
754        self.set_shape(shape)?;
755        Ok(self)
756    }
757
758    pub fn insert_axis(&mut self, axis: usize) -> TractResult<()> {
759        self.shape.insert(axis, 1);
760        self.strides.insert(axis, self.strides.get(axis).copied().unwrap_or(1));
761        Ok(())
762    }
763
764    pub fn remove_axis(&mut self, axis: usize) -> TractResult<()> {
765        ensure!(self.shape[axis] == 1, "Remove a non-1 axis: axis {} in {:?}", axis, self);
766        self.shape.remove(axis);
767        self.strides.remove(axis);
768        Ok(())
769    }
770
771    pub fn broadcast_into_rank(mut self, rank: usize) -> TractResult<Tensor> {
772        self.broadcast_to_rank(rank)?;
773        self.update_strides_and_len();
774        Ok(self)
775    }
776
777    pub fn broadcast_to_rank(&mut self, rank: usize) -> TractResult<()> {
778        if rank < self.rank() {
779            bail!("Can only broadcast to higher rank")
780        }
781        while self.shape.len() < rank {
782            self.shape.insert(0, 1)
783        }
784        self.update_strides_and_len();
785        Ok(())
786    }
787
788    pub fn broadcast_scalar_to_shape(&self, shape: &[usize]) -> TractResult<Tensor> {
789        if self.rank() > 0 {
790            bail!("broadcast_scalar_to_shape called on {:?}, which is not a salar", self);
791        }
792        unsafe fn make<T: Datum>(src: &Tensor, dst: &mut Tensor) {
793            unsafe {
794                let value: &T = src.to_scalar_unchecked::<T>();
795                dst.as_slice_mut_unchecked::<T>().iter_mut().for_each(|item| *item = value.clone())
796            };
797        }
798        unsafe {
799            let mut t = Tensor::uninitialized_dt(self.datum_type(), shape)?;
800            dispatch_datum_by_size!(make(self.datum_type())(self, &mut t));
801            Ok(t)
802        }
803    }
804
805    fn broadcast_to_shape_t<T: Datum>(&self, shape: &[usize]) -> TractResult<Tensor> {
806        unsafe {
807            let view = self.to_array_view_unchecked::<T>();
808            let mut output = view
809                .broadcast(shape)
810                .with_context(|| format!("Broadcasting {view:?} to {shape:?}"))?
811                .into_owned()
812                .into_tensor();
813            output.set_datum_type(self.datum_type());
814            Ok(output)
815        }
816    }
817
818    pub fn broadcast_to_shape(&self, shape: &[usize]) -> TractResult<Tensor> {
819        if !self.dt.is_copy() {
820            return dispatch_datum!(Self::broadcast_to_shape_t(self.dt)(self, shape));
821        }
822        ensure!(
823            self.rank() <= shape.len(),
824            "Broadcasting {self:?} to {shape:?} would lose {} axes",
825            self.rank() - shape.len()
826        );
827        let offset = shape.len() - self.rank();
828        let mut src: TVec<usize> = tvec!(1; shape.len());
829        src[offset..].copy_from_slice(self.shape());
830        ensure!(
831            izip!(&src, shape).all(|(s, d)| *s == 1 || s == d),
832            "Broadcasting {self:?} to {shape:?}"
833        );
834        // The axes from the innermost one down to the first broadcast axis are
835        // one contiguous run of the source, so only the axes above it need a
836        // coordinate walk, with a null source stride wherever they broadcast.
837        let mut split = shape.len();
838        while split > 0 && src[split - 1] == shape[split - 1] {
839            split -= 1;
840        }
841        let dt_size = self.dt.size_of();
842        let run = shape[split..].iter().product::<usize>() * dt_size;
843        let outer: usize = shape[..split].iter().product();
844        let mut src_strides: TVec<usize> = tvec!(0; split);
845        let mut acc = run;
846        for ax in (0..split).rev() {
847            src_strides[ax] = if src[ax] == 1 { 0 } else { acc };
848            acc *= src[ax];
849        }
850        let mut output = unsafe { Tensor::uninitialized_dt(self.dt, shape)? };
851        if run == 0 || outer == 0 {
852            return Ok(output);
853        }
854        let source = self.as_bytes();
855        let dst = output.as_bytes_mut();
856        let mut coords: TVec<usize> = tvec!(0; split);
857        for block in 0..outer {
858            let from: usize = izip!(&coords, &src_strides).map(|(c, s)| c * s).sum();
859            dst[block * run..][..run].copy_from_slice(&source[from..][..run]);
860            for ax in (0..split).rev() {
861                coords[ax] += 1;
862                if coords[ax] < shape[ax] {
863                    break;
864                }
865                coords[ax] = 0;
866            }
867        }
868        Ok(output)
869    }
870
871    pub fn broadcast_vector_to_shape(&self, shape: &[usize], axis: usize) -> TractResult<Tensor> {
872        ensure!(self.rank() == 1);
873        ensure!(shape[axis] == self.len());
874        if !self.datum_type().is_copy() {
875            let mut vec_shape = vec![1; shape.len()];
876            vec_shape[axis] = self.len();
877            return self.clone().into_shape(&vec_shape)?.broadcast_to_shape(shape);
878        }
879        unsafe {
880            let mut output = Tensor::uninitialized_dt(self.datum_type(), shape)?;
881            if output.len() == 0 {
882                return Ok(output);
883            }
884            let inner_len = shape[axis + 1..].iter().product::<usize>();
885
886            unsafe fn splat<T>(input: &Tensor, output: &mut Tensor, inner_len: usize)
887            where
888                T: Datum + Copy,
889            {
890                unsafe {
891                    for ix in 0..input.len() {
892                        let value: T = input.as_slice_unchecked()[ix];
893                        output.as_slice_mut_unchecked::<T>()[ix * inner_len..(ix + 1) * inner_len]
894                            .iter_mut()
895                            .for_each(|item| *item = value);
896                    }
897                }
898            }
899            dispatch_copy_by_size!(splat(self.datum_type())(&self, &mut output, inner_len));
900
901            let outer_len = shape[0..axis].iter().product::<usize>();
902            let repeat_bytes_len = inner_len * self.as_bytes().len();
903            let bytes = output.as_bytes_mut();
904            for ix in 1..outer_len {
905                bytes.copy_within(0..repeat_bytes_len, ix * repeat_bytes_len);
906            }
907
908            Ok(output)
909        }
910    }
911    pub fn assign_slice(
912        &mut self,
913        range: impl std::ops::RangeBounds<usize>,
914        src: &Tensor,
915        src_range: impl std::ops::RangeBounds<usize>,
916        axis: usize,
917    ) -> TractResult<()> {
918        self.assign_slice_at_prefix(&[], range, src, &[], src_range, axis)
919    }
920
921    /// Assign `src`'s `src_range` along `axis` into `range` along `axis`, each
922    /// taken in the sub-tensor at its prefix. A prefix indexes the leading axes
923    /// as [`Tensor::view_at_prefix`] does, and both prefixes cover the same
924    /// axes, whose extents may then differ. `axis` stays an axis of the whole
925    /// tensors, so it has to sit past the prefixes.
926    pub fn assign_slice_at_prefix(
927        &mut self,
928        prefix: &[usize],
929        range: impl std::ops::RangeBounds<usize>,
930        src: &Tensor,
931        src_prefix: &[usize],
932        src_range: impl std::ops::RangeBounds<usize>,
933        axis: usize,
934    ) -> TractResult<()> {
935        ensure!(self.rank() == src.rank());
936        ensure!(axis < self.rank());
937        let range = clip_range_bounds(self.shape[axis], range);
938        let src_range = clip_range_bounds(src.shape[axis], src_range);
939        ensure!(
940            src.datum_type() == self.datum_type(),
941            "Attempt to assign into {:?} from {:?}, datum type mismatch",
942            self.datum_type(),
943            src.datum_type()
944        );
945        ensure!(
946            src_range.len() == range.len(),
947            "Attempt to assign a range of {:?} from a range of {:?}",
948            range,
949            src_range,
950        );
951        ensure!(
952            prefix.len() == src_prefix.len() && prefix.len() <= axis,
953            "Attempt to assign axis {axis} at prefixes {prefix:?} and {src_prefix:?}"
954        );
955        ensure!(
956            izip!(prefix, self.shape()).all(|(ix, dim)| ix < dim)
957                && izip!(src_prefix, src.shape()).all(|(ix, dim)| ix < dim),
958            "Attempt to assign into {self:?} at {prefix:?} from {src:?} at {src_prefix:?}"
959        );
960        ensure!(
961            izip!(prefix.len().., &self.shape[prefix.len()..], &src.shape[prefix.len()..])
962                .all(|(ix, dst, src)| ix == axis || src == dst),
963            "Attempt to assign a {}-axis range of {:?} from a range of {:?}",
964            axis,
965            self,
966            src
967        );
968        ensure!(
969            src_range.end <= src.shape()[axis],
970            "Assigning from invalid slice (axis {}, {:?}) of {:?}",
971            axis,
972            src_range,
973            src
974        );
975        ensure!(
976            range.end <= self.shape()[axis],
977            "Assigning to invalid slice (axis {}, {:?}) of {:?}",
978            axis,
979            range,
980            self
981        );
982        unsafe { self.assign_slice_from_resolved(prefix, range, src, src_prefix, src_range, axis) };
983        Ok(())
984    }
985
986    pub unsafe fn assign_slice_unchecked(
987        &mut self,
988        range: impl std::ops::RangeBounds<usize>,
989        src: &Tensor,
990        src_range: impl std::ops::RangeBounds<usize>,
991        axis: usize,
992    ) {
993        let range = clip_range_bounds(self.shape[axis], range);
994        let src_range = clip_range_bounds(src.shape[axis], src_range);
995        unsafe { self.assign_slice_from_resolved(&[], range, src, &[], src_range, axis) };
996    }
997
998    /// The byte offset of the sub-tensor at `prefix`, which indexes the leading
999    /// axes as [`Tensor::view_at_prefix`] does.
1000    fn prefix_offset(&self, prefix: &[usize]) -> usize {
1001        izip!(prefix, &self.strides).map(|(ix, stride)| ix * *stride as usize).sum::<usize>()
1002            * self.datum_type().size_of()
1003    }
1004
1005    #[allow(clippy::ptr_eq)]
1006    unsafe fn assign_slice_from_resolved(
1007        &mut self,
1008        prefix: &[usize],
1009        range: std::ops::Range<usize>,
1010        src: &Tensor,
1011        src_prefix: &[usize],
1012        src_range: std::ops::Range<usize>,
1013        axis: usize,
1014    ) {
1015        unsafe {
1016            use ndarray::Slice;
1017            unsafe fn assign_slice_t<T: Datum>(
1018                to: &mut Tensor,
1019                to_prefix: &[usize],
1020                to_range: Range<usize>,
1021                from: &Tensor,
1022                from_prefix: &[usize],
1023                from_range: Range<usize>,
1024                axis: usize,
1025            ) {
1026                unsafe {
1027                    let mut to_view = to.to_array_view_mut_unchecked::<T>();
1028                    let mut from_view = from.to_array_view_unchecked::<T>();
1029                    for (ax, (to, from)) in izip!(to_prefix, from_prefix).enumerate() {
1030                        to_view.slice_axis_inplace(Axis(ax), Slice::from(*to..*to + 1));
1031                        from_view.slice_axis_inplace(Axis(ax), Slice::from(*from..*from + 1));
1032                    }
1033                    to_view
1034                        .slice_axis_mut(Axis(axis), Slice::from(to_range))
1035                        .assign(&from_view.slice_axis(Axis(axis), Slice::from(from_range)))
1036                }
1037            }
1038            if self.datum_type().is_copy() {
1039                // Tensors carry natural strides, so a range along `axis` is one
1040                // contiguous run per coordinate of the axes between the prefix
1041                // and it, and both sides share the run length and the trailing
1042                // block.
1043                let post = self.strides[axis] as usize * self.datum_type().size_of();
1044                let len = post * range.len();
1045                if len > 0 {
1046                    let outer: usize = self.shape[prefix.len()..axis].iter().product();
1047                    let dst_block = post * self.shape[axis];
1048                    let src_block = post * src.shape[axis];
1049                    let src_ptr = src
1050                        .plain_storage()
1051                        .as_ptr()
1052                        .add(src.prefix_offset(src_prefix) + post * src_range.start);
1053                    let aliasing = self.plain_storage().as_ptr() == src.plain_storage().as_ptr();
1054                    let dst_offset = self.prefix_offset(prefix) + post * range.start;
1055                    let dst_ptr = self.plain_storage_mut().as_mut_ptr().add(dst_offset);
1056                    for run in 0..outer {
1057                        let from = src_ptr.add(run * src_block);
1058                        let to = dst_ptr.add(run * dst_block);
1059                        if aliasing {
1060                            std::ptr::copy(from, to, len);
1061                        } else {
1062                            std::ptr::copy_nonoverlapping(from, to, len);
1063                        }
1064                    }
1065                }
1066            } else {
1067                dispatch_datum!(assign_slice_t(self.datum_type())(
1068                    self, prefix, range, src, src_prefix, src_range, axis
1069                ));
1070            }
1071        }
1072    }
1073    /// Fill `range` along `axis` with `value`, a one-element tensor of this
1074    /// tensor's datum type.
1075    pub fn fill_slice(
1076        &mut self,
1077        range: impl std::ops::RangeBounds<usize>,
1078        value: &Tensor,
1079        axis: usize,
1080    ) -> TractResult<()> {
1081        self.fill_slice_at_prefix(&[], range, value, axis)
1082    }
1083
1084    /// Fill `range` along `axis` of the sub-tensor at `prefix`, which indexes
1085    /// the leading axes as [`Tensor::view_at_prefix`] does, with `value`, a
1086    /// one-element tensor of this tensor's datum type. `axis` stays an axis of
1087    /// the whole tensor, so it has to sit past `prefix`.
1088    pub fn fill_slice_at_prefix(
1089        &mut self,
1090        prefix: &[usize],
1091        range: impl std::ops::RangeBounds<usize>,
1092        value: &Tensor,
1093        axis: usize,
1094    ) -> TractResult<()> {
1095        ensure!(axis < self.rank(), "Filling axis {axis} of {self:?}");
1096        ensure!(
1097            prefix.len() <= axis,
1098            "Filling axis {axis} of {self:?} at prefix {prefix:?}, which reaches it"
1099        );
1100        ensure!(
1101            izip!(prefix, self.shape()).all(|(ix, dim)| ix < dim),
1102            "Filling {self:?} at prefix {prefix:?}"
1103        );
1104        ensure!(
1105            value.datum_type() == self.datum_type() && value.len() == 1,
1106            "Filling {:?} with {value:?}",
1107            self.datum_type()
1108        );
1109        let range = clip_range_bounds(self.shape[axis], range);
1110        ensure!(
1111            range.end <= self.shape[axis],
1112            "Filling invalid slice (axis {axis}, {range:?}) of {self:?}"
1113        );
1114        if !self.datum_type().is_copy() {
1115            return dispatch_datum!(Self::fill_slice_t(self.datum_type())(
1116                self, prefix, range, value, axis
1117            ));
1118        }
1119        // Tensors carry natural strides, so the range is one contiguous run per
1120        // coordinate of the axes between the prefix and `axis`, and a run of one
1121        // datum grows to its whole length in log2(len) copies of itself.
1122        let dt_size = self.datum_type().size_of();
1123        let post = self.strides[axis] as usize * dt_size;
1124        let len = post * range.len();
1125        if len == 0 {
1126            return Ok(());
1127        }
1128        let block = post * self.shape[axis];
1129        let runs: usize = self.shape[prefix.len()..axis].iter().product();
1130        let start = self.prefix_offset(prefix) + range.start * post;
1131        let value = &value.as_bytes()[..dt_size];
1132        let data = self.as_bytes_mut();
1133        for run in 0..runs {
1134            let run = &mut data[start + run * block..start + run * block + len];
1135            run[..dt_size].copy_from_slice(value);
1136            let mut written = dt_size;
1137            while written < len {
1138                let grow = written.min(len - written);
1139                run.copy_within(0..grow, written);
1140                written += grow;
1141            }
1142        }
1143        Ok(())
1144    }
1145
1146    fn fill_slice_t<T: Datum>(
1147        &mut self,
1148        prefix: &[usize],
1149        range: Range<usize>,
1150        value: &Tensor,
1151        axis: usize,
1152    ) -> TractResult<()> {
1153        let value = value.try_as_plain()?.to_scalar::<T>()?.clone();
1154        let mut view = self.to_plain_array_view_mut::<T>()?;
1155        for (ax, ix) in prefix.iter().enumerate() {
1156            view.slice_axis_inplace(Axis(ax), (*ix..*ix + 1).into());
1157        }
1158        view.slice_axis_mut(Axis(axis), range.into()).fill(value);
1159        Ok(())
1160    }
1161
1162    /// Get the datum type of the tensor.
1163    #[inline]
1164    pub fn datum_type(&self) -> DatumType {
1165        self.dt
1166    }
1167
1168    /// Set the datum type of the tensor.
1169    #[inline]
1170    pub unsafe fn set_datum_type(&mut self, dt: DatumType) {
1171        self.dt = dt
1172    }
1173
1174    /// Dump the tensor in a human readable form.
1175    ///
1176    /// `force_full` will force the tensor to be dump in full even if it is big.
1177    pub fn dump(&self, force_full: bool) -> TractResult<String> {
1178        if self.is_exotic() {
1179            return Ok(format!(
1180                "{},{:?} (non-plain storage)",
1181                self.shape.iter().join(","),
1182                self.dt,
1183            ));
1184        }
1185        unsafe fn dump_t<D: Datum>(tensor: &Tensor, n: usize) -> String {
1186            unsafe {
1187                if let Some(qp) = tensor.datum_type().qparams() {
1188                    let integers = tensor.cast_to::<i32>().unwrap();
1189                    integers.as_slice_unchecked::<i32>()[0..n]
1190                        .iter()
1191                        .map(|x| format!("[{}]({})", x, qp.dq(*x)))
1192                        .join(", ")
1193                } else {
1194                    tensor.as_slice_unchecked::<D>()[0..n].iter().join(", ")
1195                }
1196            }
1197        }
1198        unsafe {
1199            let trunc = self.len() > 12 && !force_full;
1200            let data = dispatch_datum!(dump_t(self.datum_type())(
1201                self,
1202                if trunc { 12 } else { self.len() }
1203            ));
1204            Ok(format!(
1205                "{},{:?} {}{}",
1206                self.shape.iter().join(","),
1207                self.dt,
1208                data,
1209                if trunc { "..." } else { "" }
1210            ))
1211        }
1212    }
1213
1214    /// Compare two tensors, allowing for rounding errors.
1215    pub fn close_enough(
1216        &self,
1217        other: &Self,
1218        approx: impl Into<Approximation> + std::fmt::Debug,
1219    ) -> TractResult<()> {
1220        let approx = approx.into();
1221        if self.shape() != other.shape() {
1222            bail!("Shape mismatch {:?} != {:?}", self.shape(), other.shape())
1223        }
1224        if let Approximation::Ulp(max_ulp) = approx {
1225            return self.ulp_close_enough(other, max_ulp);
1226        }
1227        let (atol, rtol, outliers) = approx.atol_rtol_outliers(&self.datum_type());
1228        let ma = self.cast_to::<f32>()?;
1229        let ma = ma.to_plain_array_view::<f32>()?;
1230        let mb = other.cast_to::<f32>()?;
1231        let mb = mb.to_plain_array_view::<f32>()?;
1232        let mut first_outlier = None;
1233        let mut outliers_count = 0;
1234        ndarray::indices_of(&ma).into_iter().for_each(|indices| {
1235            let a = ma[&indices];
1236            let b = mb[&indices];
1237            if !((a.is_nan() && b.is_nan())
1238                || (a.is_infinite() && b.is_infinite() && a.signum() == b.signum())
1239                || (a - b).abs() <= atol as f32 + rtol as f32 * b.abs())
1240            {
1241                if outliers_count == 0 {
1242                    first_outlier = Some(indices.as_array_view().to_vec());
1243                }
1244                outliers_count += 1;
1245            }
1246        });
1247        if self.volume() > 0 && outliers_count as f64 / self.volume() as f64 > outliers {
1248            let indices = first_outlier.unwrap();
1249            let a = ma[&*indices];
1250            let b = mb[&*indices];
1251            let ulp = self
1252                .max_ulp_distance(other)
1253                .map(|(d, _)| format!("{d}"))
1254                .unwrap_or_else(|_| "n/a".to_string());
1255            bail!(
1256                "Mismatch. First outlier: {:?} for {:?}) at {:?} {} != {}. Outliers: {} / {} = {:0.5} > {:0.5}. Max ULP ({:?}): {}.",
1257                approx,
1258                self.datum_type(),
1259                indices,
1260                a,
1261                b,
1262                outliers_count,
1263                self.volume(),
1264                outliers_count as f64 / self.volume() as f64,
1265                outliers,
1266                self.ulp_comparison_dt(),
1267                ulp,
1268            );
1269        }
1270        Ok(())
1271    }
1272
1273    /// The float type ULP distances against this tensor are measured in.
1274    ///
1275    /// Float tensors are compared in their own type, so an f16 comparison stays an
1276    /// f16 comparison. Anything else falls back to f32, matching what
1277    /// `close_enough` does for its tolerance check.
1278    pub fn ulp_comparison_dt(&self) -> DatumType {
1279        match self.datum_type() {
1280            dt @ (DatumType::F16 | DatumType::F32 | DatumType::F64) => dt,
1281            _ => DatumType::F32,
1282        }
1283    }
1284
1285    /// Largest integer ULP distance between `self` and `other`, and the flat index
1286    /// where it occurs.
1287    ///
1288    /// Comparison happens in [`Self::ulp_comparison_dt`]. See [`crate::ulp`] for
1289    /// the exact convention around signed zeros, infinities and NaN.
1290    pub fn max_ulp_distance(&self, other: &Self) -> TractResult<(u64, Option<usize>)> {
1291        if self.shape() != other.shape() {
1292            bail!("Shape mismatch {:?} != {:?}", self.shape(), other.shape())
1293        }
1294        let dt = self.ulp_comparison_dt();
1295        let a = self.cast_to_dt(dt)?;
1296        let b = other.cast_to_dt(dt)?;
1297        fn worst<D: Datum + crate::ulp::UlpFloat>(
1298            a: &Tensor,
1299            b: &Tensor,
1300        ) -> TractResult<(u64, Option<usize>)> {
1301            let a = a.to_plain_array_view::<D>()?;
1302            let b = b.to_plain_array_view::<D>()?;
1303            Ok(crate::ulp::max_ulp_distance(a.iter().copied(), b.iter().copied()))
1304        }
1305        match dt {
1306            DatumType::F16 => worst::<f16>(&a, &b),
1307            DatumType::F32 => worst::<f32>(&a, &b),
1308            DatumType::F64 => worst::<f64>(&a, &b),
1309            dt => bail!("No ULP comparison for {dt:?}"),
1310        }
1311    }
1312
1313    /// Compare two tensors by integer ULP distance, accepting a distance up to
1314    /// `max_ulp`.
1315    fn ulp_close_enough(&self, other: &Self, max_ulp: u64) -> TractResult<()> {
1316        let (worst, at) = self.max_ulp_distance(other)?;
1317        if worst <= max_ulp {
1318            return Ok(());
1319        }
1320        let dt = self.ulp_comparison_dt();
1321        let indices = at
1322            .map(|flat| {
1323                let mut rest = flat;
1324                let mut indices = vec![0; self.rank()];
1325                for (ix, dim) in self.shape().iter().enumerate().rev() {
1326                    indices[ix] = rest % dim;
1327                    rest /= dim;
1328                }
1329                format!("{indices:?}")
1330            })
1331            .unwrap_or_else(|| "?".to_string());
1332        let a = self.cast_to::<f64>()?;
1333        let b = other.cast_to::<f64>()?;
1334        let (a, b) = (a.to_plain_array_view::<f64>()?, b.to_plain_array_view::<f64>()?);
1335        let flat = at.unwrap_or(0);
1336        bail!(
1337            "Mismatch. Max ULP distance ({dt:?}): {} > {}, at {} ({} != {}).",
1338            worst,
1339            max_ulp,
1340            indices,
1341            a.iter().nth(flat).copied().unwrap_or(f64::NAN),
1342            b.iter().nth(flat).copied().unwrap_or(f64::NAN),
1343        );
1344    }
1345
1346    /// Transform the tensor into a `ndarray::Array`.
1347    pub fn into_plain_array<D: Datum>(self) -> TractResult<ArrayD<D>> {
1348        Ok(self.to_plain_array_view::<D>()?.to_owned())
1349    }
1350
1351    /// Transform the tensor into a `ndarray::Array`.
1352    pub unsafe fn into_array_unchecked<D: Datum>(self) -> ArrayD<D> {
1353        unsafe { self.to_array_view_unchecked::<D>().to_owned() }
1354    }
1355
1356    /// Returns a plain array view of the tensor.
1357    ///
1358    /// Errors if the storage is not plain or the datum type does not match `D`.
1359    #[inline]
1360    pub fn to_plain_array_view<D: Datum>(&self) -> TractResult<ArrayViewD<'_, D>> {
1361        self.try_as_plain()?.to_array_view::<D>()
1362    }
1363
1364    /// Returns a mutable plain array view of the tensor.
1365    ///
1366    /// Errors if the storage is not plain or the datum type does not match `D`.
1367    #[inline]
1368    pub fn to_plain_array_view_mut<D: Datum>(&mut self) -> TractResult<ArrayViewMutD<'_, D>> {
1369        self.check_for_access::<D>()?;
1370        ensure!(self.storage.as_plain_mut().is_some(), "Tensor storage is not plain");
1371        unsafe { Ok(self.to_array_view_mut_unchecked()) }
1372    }
1373
1374    fn check_for_access<D: Datum>(&self) -> TractResult<()> {
1375        ensure!(
1376            self.datum_type().unquantized() == D::datum_type().unquantized(),
1377            "Tensor datum type error: tensor is {:?}, accessed as {:?}",
1378            self.datum_type(),
1379            D::datum_type(),
1380        );
1381        Ok(())
1382    }
1383
1384    /// Transform the data as a `ndarray::Array`.
1385    pub unsafe fn to_array_view_unchecked<D: Datum>(&self) -> ArrayViewD<'_, D> {
1386        if self.len() != 0 {
1387            unsafe {
1388                ArrayViewD::from_shape_ptr(&*self.shape, self.plain_storage().as_ptr() as *const D)
1389            }
1390        } else {
1391            ArrayViewD::from_shape(&*self.shape, &[]).unwrap()
1392        }
1393    }
1394
1395    /// Transform the data as a mutable `ndarray::Array`.
1396    pub unsafe fn to_array_view_mut_unchecked<D: Datum>(&mut self) -> ArrayViewMutD<'_, D> {
1397        if self.len() != 0 {
1398            unsafe {
1399                let ptr = self.plain_storage_mut().as_mut_ptr() as *mut D;
1400                ArrayViewMutD::from_shape_ptr(&*self.shape, ptr)
1401            }
1402        } else {
1403            ArrayViewMutD::from_shape(&*self.shape, &mut []).unwrap()
1404        }
1405    }
1406
1407    /// Access the data as a pointer.
1408    pub fn as_ptr<D: Datum>(&self) -> TractResult<*const D> {
1409        self.check_for_access::<D>()?;
1410        Ok(self.plain_storage().as_ptr() as *const D)
1411    }
1412
1413    /// Access the data as a pointer.
1414    pub unsafe fn as_ptr_unchecked<D: Datum>(&self) -> *const D {
1415        self.plain_storage().as_ptr() as *const D
1416    }
1417
1418    /// Access the data as a pointer.
1419    pub unsafe fn as_ptr_mut_unchecked<D: Datum>(&mut self) -> *mut D {
1420        self.plain_storage_mut().as_mut_ptr() as *mut D
1421    }
1422
1423    /// Access the data as a mutable pointer.
1424    pub fn as_ptr_mut<D: Datum>(&mut self) -> TractResult<*mut D> {
1425        self.as_ptr::<D>().map(|p| p as *mut D)
1426    }
1427
1428    /// Access the data as a slice.
1429    pub unsafe fn as_slice_unchecked<D: Datum>(&self) -> &[D] {
1430        if self.storage.byte_len() == 0 {
1431            &[]
1432        } else {
1433            unsafe { std::slice::from_raw_parts::<D>(self.as_ptr_unchecked(), self.len()) }
1434        }
1435    }
1436
1437    /// Access the data as a mutable slice.
1438    pub unsafe fn as_slice_mut_unchecked<D: Datum>(&mut self) -> &mut [D] {
1439        if self.storage.byte_len() == 0 {
1440            &mut []
1441        } else {
1442            unsafe { std::slice::from_raw_parts_mut::<D>(self.as_ptr_mut_unchecked(), self.len()) }
1443        }
1444    }
1445
1446    /// Make the tensor a scalar tensor (assumes it contains a single value).
1447    pub fn to_scalar_tensor(&self) -> TractResult<Tensor> {
1448        fn to_scalar_tensor_t<D: Datum>(t: &Tensor) -> TractResult<Tensor> {
1449            Ok(litteral::tensor0(t.try_as_plain()?.to_scalar::<D>()?.clone()))
1450        }
1451        dispatch_datum!(to_scalar_tensor_t(self.datum_type())(self))
1452    }
1453
1454    /// Access the data as a scalar.
1455    pub unsafe fn to_scalar_unchecked<D: Datum>(&self) -> &D {
1456        unsafe { &*(self.plain_storage().as_ptr() as *const D) }
1457    }
1458
1459    /// Mutable access the data as a scalar.
1460    pub fn to_scalar_mut<D: Datum>(&mut self) -> TractResult<&mut D> {
1461        self.check_for_access::<D>()?;
1462        if self.len() == 0 {
1463            bail!("to_scalar_mut called on empty tensor ({:?})", self)
1464        }
1465        if self.len() > 1 {
1466            bail!("to_scalar called on a tensor with multiple values ({:?})", self)
1467        }
1468        unsafe { Ok(self.to_scalar_mut_unchecked()) }
1469    }
1470
1471    /// Mutable access the data as a scalar.
1472    pub unsafe fn to_scalar_mut_unchecked<D: Datum>(&mut self) -> &mut D {
1473        unsafe { &mut *(self.plain_storage_mut().as_mut_ptr() as *mut D) }
1474    }
1475
1476    pub fn as_bytes(&self) -> &[u8] {
1477        self.plain_storage().as_bytes()
1478    }
1479
1480    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
1481        self.plain_storage_mut().as_bytes_mut()
1482    }
1483
1484    unsafe fn is_uniform_t<T: Datum>(&self) -> bool {
1485        let slice = unsafe { self.as_slice_unchecked::<T>() };
1486        slice[1..].iter().all(|x| x == &slice[0])
1487    }
1488
1489    pub fn is_uniform(&self) -> bool {
1490        if self.is_exotic() {
1491            return false;
1492        }
1493        if self.len() <= 1 {
1494            return true;
1495        }
1496        unsafe { dispatch_datum!(Tensor::is_uniform_t(self.datum_type())(self)) }
1497    }
1498
1499    unsafe fn as_uniform_t<T: Datum>(&self) -> Tensor {
1500        let v: T = unsafe { self.as_slice_unchecked::<T>() }[0].clone();
1501        litteral::tensor0(v)
1502    }
1503
1504    pub fn as_uniform(&self) -> Option<Tensor> {
1505        if self.len() >= 1 && self.is_uniform() {
1506            unsafe {
1507                let mut t = dispatch_datum!(Tensor::as_uniform_t(self.datum_type())(self));
1508                t.set_datum_type(self.datum_type());
1509                Some(t)
1510            }
1511        } else {
1512            None
1513        }
1514    }
1515
1516    pub fn is_all_zero(&self) -> TractResult<bool> {
1517        Ok(self.len() == 0 || self.as_uniform().map(|t| t.is_zero().unwrap()).unwrap_or(false))
1518    }
1519
1520    pub fn is_zero(&self) -> TractResult<bool> {
1521        Ok(self == &Tensor::zero_scalar_dt(self.dt)?)
1522    }
1523
1524    unsafe fn natural_cast<
1525        Source: Datum + num_traits::AsPrimitive<Target>,
1526        Target: Datum + Copy,
1527    >(
1528        &self,
1529        other: &mut Tensor,
1530    ) {
1531        unsafe {
1532            self.as_slice_unchecked::<Source>()
1533                .iter()
1534                .zip(other.as_slice_mut_unchecked::<Target>().iter_mut())
1535                .for_each(|(s, d)| *d = s.as_())
1536        };
1537    }
1538
1539    unsafe fn cast_number_to_bool<Source: Datum + num_traits::Zero>(&self, other: &mut Tensor) {
1540        unsafe {
1541            self.as_slice_unchecked::<Source>()
1542                .iter()
1543                .zip(other.as_slice_mut_unchecked::<bool>().iter_mut())
1544                .for_each(|(s, d)| *d = !s.is_zero());
1545        }
1546    }
1547
1548    unsafe fn cast_from_string<Target: Datum + core::str::FromStr>(
1549        &self,
1550        other: &mut Tensor,
1551    ) -> TractResult<()> {
1552        unsafe {
1553            for (s, d) in self
1554                .as_slice_unchecked::<String>()
1555                .iter()
1556                .zip(other.as_slice_mut_unchecked::<Target>().iter_mut())
1557            {
1558                *d = s
1559                    .parse()
1560                    .map_err(|_| format_err!("Can not parse as {:?}", Target::datum_type()))?;
1561            }
1562            Ok(())
1563        }
1564    }
1565
1566    unsafe fn cast_to_string<Source: Datum>(&self, other: &mut Tensor) {
1567        unsafe {
1568            for (s, d) in self
1569                .as_slice_unchecked::<Source>()
1570                .iter()
1571                .zip(other.as_slice_mut_unchecked::<String>().iter_mut())
1572            {
1573                *d = s.to_string()
1574            }
1575        }
1576    }
1577
1578    /// Optionnaly convert data to a tensor for a new DatumType.
1579    pub fn cast_to<D: Datum>(&self) -> TractResult<Cow<'_, Tensor>> {
1580        self.cast_to_dt(D::datum_type())
1581    }
1582
1583    /// Optionnaly convert data to a tensor for a new DatumType.
1584    #[allow(clippy::redundant_closure_call)]
1585    pub fn cast_to_dt(&self, dst_dt: DatumType) -> TractResult<Cow<'_, Tensor>> {
1586        unsafe {
1587            if self.dt == dst_dt {
1588                return Ok(Cow::Borrowed(self));
1589            }
1590            if self.dt == TDim::datum_type() && (dst_dt.is_integer() || dst_dt.is_float()) {
1591                let slice = self.as_slice_unchecked::<TDim>();
1592                let mut ints = Self::uninitialized::<i64>(&self.shape)?;
1593                let ints_slice = ints.as_slice_mut_unchecked::<i64>();
1594                for i in 0..self.len() {
1595                    ints_slice[i] = slice[i].to_i64()?;
1596                }
1597                return Ok(Cow::Owned(ints.cast_to_dt(dst_dt)?.into_owned()));
1598            }
1599            if self.dt == bool::datum_type()
1600                && (dst_dt.is_integer() || dst_dt.is_float() || dst_dt == TDim::datum_type())
1601            {
1602                let slice = self.as_slice_unchecked::<bool>();
1603                let mut ints = Self::uninitialized::<i8>(&self.shape)?;
1604                let ints_slice = ints.as_slice_mut_unchecked::<i8>();
1605                for i in 0..self.len() {
1606                    ints_slice[i] = slice[i] as usize as i8;
1607                }
1608                return Ok(Cow::Owned(ints.cast_to_dt(dst_dt)?.into_owned()));
1609            }
1610            let mut result = Self::uninitialized_dt(dst_dt, &self.shape)?;
1611            if self.dt == DatumType::String {
1612                dispatch_numbers!(Self::cast_from_string(dst_dt)(self, &mut result))?;
1613                return Ok(Cow::Owned(result));
1614            }
1615            if dst_dt == DatumType::String {
1616                dispatch_datum!(Self::cast_to_string(self.dt)(self, &mut result));
1617                return Ok(Cow::Owned(result));
1618            }
1619            macro_rules! n {
1620                ($source:ty) => {
1621                    if <$source>::datum_type() == self.datum_type() {
1622                        match dst_dt {
1623                            DatumType::I8 => self.natural_cast::<$source, i8>(&mut result),
1624                            DatumType::I16 => self.natural_cast::<$source, i16>(&mut result),
1625                            DatumType::I32 => self.natural_cast::<$source, i32>(&mut result),
1626                            DatumType::I64 => self.natural_cast::<$source, i64>(&mut result),
1627                            DatumType::U8 => self.natural_cast::<$source, u8>(&mut result),
1628                            DatumType::U16 => self.natural_cast::<$source, u16>(&mut result),
1629                            DatumType::U32 => self.natural_cast::<$source, u32>(&mut result),
1630                            DatumType::U64 => self.natural_cast::<$source, u64>(&mut result),
1631                            DatumType::F16 => self.natural_cast::<$source, f16>(&mut result),
1632                            DatumType::F32 => self.natural_cast::<$source, f32>(&mut result),
1633                            DatumType::F64 => self.natural_cast::<$source, f64>(&mut result),
1634                            DatumType::TDim => {
1635                                let ints = self.cast_to::<i32>()?;
1636                                let slice = ints.as_slice_unchecked::<i32>();
1637                                let result = result.as_slice_mut_unchecked::<TDim>();
1638                                for i in 0..self.len() {
1639                                    result[i] = slice[i].into();
1640                                }
1641                            }
1642                            DatumType::Bool => self.cast_number_to_bool::<$source>(&mut result),
1643                            _ => todo!(),
1644                        }
1645                        return Ok(Cow::Owned(result));
1646                    };
1647                };
1648            }
1649            //If there is no quantization
1650            if !dst_dt.is_quantized() && !self.datum_type().is_quantized() {
1651                n!(u8);
1652                n!(u16);
1653                n!(u32);
1654                n!(u64);
1655                n!(i8);
1656                n!(i16);
1657                n!(i32);
1658                n!(i64);
1659                n!(f16);
1660                n!(f32);
1661                n!(f64);
1662            } else {
1663                let (s_zp, s_scale) = self.datum_type().zp_scale();
1664                let (d_zp, d_scale) = dst_dt.zp_scale();
1665                if self.datum_type().is_quantized() && dst_dt.is_float() {
1666                    macro_rules! q_to_fp {
1667                        ($source:ty, $dest:ty) => {
1668                            if <$source>::datum_type().unquantized()
1669                                == self.datum_type().unquantized()
1670                                && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1671                            {
1672                                self.as_slice_unchecked::<$source>()
1673                                    .iter()
1674                                    .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1675                                    .for_each(|(&s, d)| {
1676                                        *d = (s as $dest - s_zp as $dest) * s_scale as $dest;
1677                                    });
1678                                return Ok(Cow::Owned(result));
1679                            }
1680                        };
1681                    }
1682                    q_to_fp!(i8, f64);
1683                    q_to_fp!(i8, f32);
1684                    q_to_fp!(u8, f64);
1685                    q_to_fp!(u8, f32);
1686                }
1687                //TODO: optimize scale_by
1688                macro_rules! q8_to_q8 {
1689                    ($typ:ty) => {
1690                        if dst_dt.unquantized() == <$typ>::datum_type() {
1691                            self.as_slice_unchecked::<$typ>()
1692                                .iter()
1693                                .zip(result.as_slice_mut_unchecked::<$typ>().iter_mut())
1694                                .for_each(|(&s, d)| {
1695                                    *d = (d_zp as i32
1696                                        + scale_by(s as i32 - s_zp as i32, s_scale / d_scale))
1697                                    .clamp_cast()
1698                                });
1699                            return Ok(Cow::Owned(result));
1700                        }
1701                    };
1702                }
1703
1704                macro_rules! q_via_f32 {
1705                    ($source:ty, $dest:ty, $round:expr) => {
1706                        if <$source>::datum_type().unquantized() == self.datum_type().unquantized()
1707                            && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1708                        {
1709                            self.as_slice_unchecked::<$source>()
1710                                .iter()
1711                                .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1712                                .for_each(|(&s, d)| {
1713                                    let s_float = (s as f32 - s_zp as f32) * s_scale as f32;
1714                                    let d_float = s_float as f32 / d_scale as f32 + d_zp as f32;
1715                                    *d = $round(d_float);
1716                                });
1717                            return Ok(Cow::Owned(result));
1718                        }
1719                    };
1720                }
1721
1722                macro_rules! q_n {
1723                    (clamp $source:ty, $dest:ty) => {{
1724                        if <$source>::datum_type().unquantized() == self.datum_type().unquantized()
1725                            && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1726                        {
1727                            self.as_slice_unchecked::<$source>()
1728                                .iter()
1729                                .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1730                                .for_each(|(&s, d)| {
1731                                    *d = s.clamp_cast();
1732                                });
1733                            return Ok(Cow::Owned(result));
1734                        }
1735                    }};
1736                    ($source:ty, $dest:ty) => {{
1737                        if <$source>::datum_type().unquantized() == self.datum_type().unquantized()
1738                            && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1739                        {
1740                            self.as_slice_unchecked::<$source>()
1741                                .iter()
1742                                .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1743                                .for_each(|(&s, d)| {
1744                                    *d = s as $dest;
1745                                });
1746                            return Ok(Cow::Owned(result));
1747                        }
1748                    }};
1749                }
1750
1751                if dst_dt.unquantized() == self.datum_type().unquantized()
1752                    && dst_dt.is_quantized()
1753                    && self.datum_type().is_quantized()
1754                {
1755                    q8_to_q8!(i8);
1756                    q8_to_q8!(u8);
1757                }
1758
1759                q_via_f32!(f32, i8, |f| round_ties_to_even(f).clamp_cast());
1760                q_via_f32!(f32, u8, |f| round_ties_to_even(f).clamp_cast());
1761                q_via_f32!(f32, i32, |f| round_ties_to_even(f).clamp_cast());
1762                q_via_f32!(i8, f32, |f| f);
1763                q_via_f32!(u8, f32, |f| f);
1764                q_via_f32!(i32, f32, |f| f);
1765
1766                if dst_dt.is_quantized() && self.datum_type().is_quantized() {
1767                    q_via_f32!(u8, i8, |f| round_ties_to_even(f).clamp_cast());
1768                    q_via_f32!(i8, u8, |f| round_ties_to_even(f).clamp_cast());
1769                    q_via_f32!(i32, u8, |f| round_ties_to_even(f).clamp_cast());
1770                    q_via_f32!(i32, i8, |f| round_ties_to_even(f).clamp_cast());
1771                    q_via_f32!(u8, i32, |f| round_ties_to_even(f).clamp_cast());
1772                    q_via_f32!(i8, i32, |f| round_ties_to_even(f).clamp_cast());
1773
1774                    // ensure cast to different scale offset work
1775                    q_via_f32!(i8, i8, |f| round_ties_to_even(f).clamp_cast());
1776                    q_via_f32!(u8, u8, |f| round_ties_to_even(f).clamp_cast());
1777                }
1778
1779                q_n!(i8, i32);
1780                q_n!(i8, u32);
1781                q_n!(u8, i32);
1782                q_n!(u8, u32);
1783                q_n!(clamp i32, i8);
1784                q_n!(clamp i32, u8);
1785                q_n!(clamp u32, i8);
1786                q_n!(clamp u32, u8);
1787                q_n!(i8, i8);
1788                q_n!(u8, u8);
1789                q_n!(i32, i32);
1790                q_n!(u32, u32);
1791            }
1792
1793            bail!("Unsupported cast from {:?} to {:?}", self.dt, dst_dt)
1794        }
1795    }
1796
1797    /// Access the data as a scalar, after a cast.
1798    pub fn cast_to_scalar<D: Datum + Copy>(&self) -> TractResult<D> {
1799        let casted = self.cast_to::<D>()?;
1800        casted.try_as_plain()?.to_scalar::<D>().copied()
1801    }
1802
1803    /// Access the nth element of the tensor, returned as a 0-rank Tensor
1804    pub fn nth(&self, nth: usize) -> TractResult<Tensor> {
1805        if nth >= self.len() {
1806            bail!(
1807                "nth called with {}th element on a tensor of len {} ({:?}",
1808                nth,
1809                self.len(),
1810                self
1811            );
1812        }
1813        unsafe fn nth_t<T: Datum>(me: &Tensor, nth: usize, output: &mut Tensor) {
1814            unsafe {
1815                let value = me.as_slice_unchecked::<T>()[nth].clone();
1816                std::ptr::write(output.as_slice_mut_unchecked::<T>().as_mut_ptr(), value);
1817            }
1818        }
1819        unsafe {
1820            let mut output = Tensor::uninitialized_dt(self.datum_type(), &[])?;
1821            dispatch_datum_by_size!(nth_t(self.datum_type())(self, nth, &mut output));
1822            Ok(output)
1823        }
1824    }
1825
1826    /// Strict equality test on tensors.
1827    fn eq_dt(&self, other: &Tensor) -> TractResult<bool> {
1828        unsafe fn eq_t<D: Datum>(me: &Tensor, other: &Tensor) -> TractResult<bool> {
1829            unsafe {
1830                if D::datum_type().is_float() {
1831                    return dispatch_floatlike!(float_eq_t(D::datum_type())(me, other));
1832                }
1833                Ok(izip!(me.as_slice_unchecked::<D>(), other.as_slice_unchecked::<D>())
1834                    .all(|(a, b)| a == b))
1835            }
1836        }
1837
1838        unsafe fn float_eq_t<D: Datum + Float>(me: &Tensor, other: &Tensor) -> TractResult<bool> {
1839            unsafe {
1840                Ok(izip!(me.as_slice_unchecked::<D>(), other.as_slice_unchecked::<D>())
1841                    .all(|(a, b)| (a.is_nan() && b.is_nan()) || a == b))
1842            }
1843        }
1844
1845        unsafe {
1846            Ok(self.datum_type() == other.datum_type()
1847                && self.shape() == other.shape()
1848                && dispatch_datum!(eq_t(self.dt)(self, other))?)
1849        }
1850    }
1851
1852    fn from_datum<T: Datum>(mut it: ArrayD<T>) -> Tensor {
1853        unsafe {
1854            let mut t = Self::uninitialized::<T>(it.shape()).unwrap();
1855            if let Some(slice) = it.as_slice_mut() {
1856                if t.datum_type().is_copy() {
1857                    std::ptr::copy_nonoverlapping(
1858                        slice.as_ptr() as *const i8,
1859                        t.as_ptr_mut_unchecked(),
1860                        t.plain_storage().layout().size(),
1861                    );
1862                } else {
1863                    t.as_slice_mut_unchecked::<T>()
1864                        .iter_mut()
1865                        .zip(slice.iter_mut())
1866                        .for_each(|(t, s)| *t = std::mem::take(s));
1867                }
1868                return t;
1869            }
1870            if it.strides().iter().all(|&s| s > 0) && it.as_slice_memory_order().is_some() {
1871                let mut len_and_strides: TVec<(usize, usize)> = tvec!();
1872                for (len, stride) in itertools::izip!(it.shape(), it.strides(), t.strides())
1873                    .sorted_by_key(|(_, src, _)| *src)
1874                    .map(|(l, _, dst)| (*l as isize, *dst))
1875                {
1876                    if !len_and_strides.is_empty()
1877                        && len_and_strides.last().unwrap().1 * len_and_strides.last().unwrap().0
1878                            == stride as usize
1879                    {
1880                        len_and_strides.last_mut().unwrap().0 *= len as usize;
1881                    } else {
1882                        len_and_strides.push((len as usize, stride as usize));
1883                    }
1884                }
1885                len_and_strides.reverse();
1886                crate::scatter::scatter_contig_data(
1887                    it.as_ptr(),
1888                    t.as_ptr_mut_unchecked(),
1889                    &len_and_strides,
1890                );
1891                return t;
1892            }
1893            // finally use ndarray into_iter()
1894            t.as_slice_mut_unchecked().iter_mut().zip(it).for_each(|(t, a)| *t = a);
1895            t
1896        }
1897    }
1898
1899    pub fn deep_clone(&self) -> Tensor {
1900        if self.is_exotic() {
1901            return Tensor {
1902                dt: self.dt,
1903                shape: self.shape.clone(),
1904                strides: self.strides.clone(),
1905                len: self.len,
1906                storage: self.storage.deep_clone(),
1907            };
1908        }
1909        unsafe {
1910            let mut tensor = Tensor::uninitialized_dt(self.datum_type(), self.shape()).unwrap();
1911            if self.len() > 0 {
1912                if self.dt.is_copy() {
1913                    self.plain_storage().as_ptr().copy_to_nonoverlapping(
1914                        tensor.as_bytes_mut().as_mut_ptr(),
1915                        self.plain_storage().layout().size(),
1916                    )
1917                } else if self.dt == DatumType::String {
1918                    tensor
1919                        .as_slice_mut_unchecked::<String>()
1920                        .clone_from_slice(self.as_slice_unchecked());
1921                } else if self.dt == DatumType::Blob {
1922                    tensor
1923                        .as_slice_mut_unchecked::<Blob>()
1924                        .clone_from_slice(self.as_slice_unchecked());
1925                } else if self.dt == DatumType::TDim {
1926                    tensor
1927                        .as_slice_mut_unchecked::<TDim>()
1928                        .clone_from_slice(self.as_slice_unchecked());
1929                }
1930            }
1931            tensor
1932        }
1933    }
1934
1935    pub fn slice(&self, axis: usize, start: usize, end: usize) -> TractResult<Tensor> {
1936        if axis >= self.rank() {
1937            bail!("Can not slice at axis {} tensor {:?}", axis, self);
1938        }
1939        if start > self.shape[axis] || end > self.shape[axis] || start >= end {
1940            bail!("Invalid slicing range {start}..{end} on axis {axis} for {self:?}");
1941        }
1942        let mut shape: TVec<usize> = self.shape().into();
1943        shape[axis] = end - start;
1944        unsafe {
1945            let mut tensor = Tensor::uninitialized_dt(self.datum_type(), &shape)?;
1946            tensor.assign_slice_from_resolved(&[], 0..end - start, self, &[], start..end, axis);
1947            Ok(tensor)
1948        }
1949    }
1950
1951    #[inline]
1952    pub fn view(&self) -> view::TensorView<'_> {
1953        unsafe { view::TensorView::view(self) }
1954    }
1955
1956    #[inline]
1957    pub fn view_at_prefix(&self, prefix: &[usize]) -> TractResult<view::TensorView<'_>> {
1958        view::TensorView::at_prefix(self, prefix)
1959    }
1960
1961    #[inline]
1962    pub fn view_offsetting(&self, coords: &[usize]) -> TractResult<view::TensorView<'_>> {
1963        view::TensorView::offsetting(self, coords)
1964    }
1965
1966    #[inline]
1967    pub unsafe fn view_offsetting_unchecked(&self, coords: &[usize]) -> view::TensorView<'_> {
1968        unsafe { view::TensorView::offsetting_unchecked(self, coords) }
1969    }
1970
1971    #[inline]
1972    pub fn view_mut(&mut self) -> view::TensorView<'_> {
1973        unsafe { view::TensorView::view(self) }
1974    }
1975
1976    #[inline]
1977    pub fn view_at_prefix_mut(&mut self, prefix: &[usize]) -> TractResult<view::TensorView<'_>> {
1978        view::TensorView::at_prefix(self, prefix)
1979    }
1980
1981    #[inline]
1982    pub fn view_offsetting_mut(&mut self, coords: &[usize]) -> TractResult<view::TensorView<'_>> {
1983        view::TensorView::offsetting(self, coords)
1984    }
1985
1986    /// Offsets the tensor as an i8 type if it's an u8 type, otherwise passes it unchanged.
1987    pub fn offset_u8_as_i8(self: &Arc<Self>) -> Arc<Self> {
1988        let mut t = if let DatumType::U8 = self.dt.unquantized() {
1989            self.try_as_plain()
1990                .unwrap()
1991                .to_array_view::<u8>()
1992                .unwrap()
1993                .mapv(|v| v.wrapping_sub(128) as i8)
1994                .into_tensor()
1995        } else {
1996            return self.clone();
1997        };
1998
1999        if let DatumType::QU8(qp) = self.dt {
2000            if let QParams::ZpScale { zero_point, scale } = qp {
2001                t.dt = DatumType::QI8(QParams::ZpScale { zero_point: zero_point - 128, scale });
2002            } else {
2003                t.dt = DatumType::QI8(qp);
2004            }
2005        }
2006
2007        t.into_arc_tensor()
2008    }
2009
2010    /// Offsets the tensor as an u8 type if it's an i8 type, otherwise passes it unchanged.
2011    pub fn offset_i8_as_u8(self: &Arc<Self>) -> Arc<Self> {
2012        let mut t = if let DatumType::I8 = self.dt.unquantized() {
2013            self.try_as_plain()
2014                .unwrap()
2015                .to_array_view::<i8>()
2016                .unwrap()
2017                .mapv(|v| (v as u8).wrapping_add(128))
2018                .into_tensor()
2019        } else {
2020            return self.clone();
2021        };
2022
2023        if let DatumType::QI8(qp) = self.dt {
2024            if let QParams::ZpScale { zero_point, scale } = qp {
2025                t.dt = DatumType::QU8(QParams::ZpScale { zero_point: zero_point + 128, scale });
2026            } else {
2027                t.dt = DatumType::QU8(qp);
2028            }
2029        }
2030        t.into_arc_tensor()
2031    }
2032
2033    pub fn to_aligned_default(&self) -> TractResult<Self> {
2034        if self.dt.is_copy() {
2035            unsafe {
2036                let mut t = Self::uninitialized_dt(self.dt, &self.shape)?;
2037                t.as_bytes_mut().copy_from_slice(self.as_bytes());
2038                Ok(t)
2039            }
2040        } else {
2041            let mut t = Self::zero_dt(self.dt, &self.shape)?;
2042            if self.dt == String::datum_type() {
2043                t.try_as_plain_mut()?
2044                    .as_slice_mut::<String>()?
2045                    .clone_from_slice(self.try_as_plain()?.as_slice()?);
2046            } else if self.dt == Blob::datum_type() {
2047                t.try_as_plain_mut()?
2048                    .as_slice_mut::<Blob>()?
2049                    .clone_from_slice(self.try_as_plain()?.as_slice()?);
2050            } else if self.dt == TDim::datum_type() {
2051                t.try_as_plain_mut()?
2052                    .as_slice_mut::<TDim>()?
2053                    .clone_from_slice(self.try_as_plain()?.as_slice()?);
2054            }
2055            Ok(t)
2056        }
2057    }
2058
2059    pub fn natural_strides(shape: &[usize]) -> TVec<isize> {
2060        let mut strides = tvec!();
2061        compute_natural_stride_to(&mut strides, shape);
2062        strides
2063    }
2064
2065    pub fn into_blob(mut self) -> TractResult<Blob> {
2066        ensure!(self.dt.is_copy());
2067        let storage =
2068            std::mem::replace(&mut self.storage, StorageKind::Plain(PlainStorage::default()));
2069        Ok(storage.into_plain().context("Storage is not plain")?.into_blob())
2070    }
2071}
2072
2073impl PartialEq for Tensor {
2074    fn eq(&self, other: &Tensor) -> bool {
2075        if self.dt != other.dt || self.shape != other.shape {
2076            return false;
2077        }
2078        match (self.storage.as_plain(), other.storage.as_plain()) {
2079            (Some(_), Some(_)) => self.eq_dt(other).unwrap_or(false),
2080            (None, None) => self.storage == other.storage,
2081            _ => false,
2082        }
2083    }
2084}
2085
2086impl Eq for Tensor {}
2087
2088impl fmt::Debug for Tensor {
2089    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
2090        let content = self.dump(false).unwrap_or_else(|e| format!("Error : {e:?}"));
2091        write!(formatter, "{content}")
2092    }
2093}
2094
2095#[cfg(feature = "complex")]
2096pub fn reinterpret_inner_dim_as_complex(mut t: Tensor) -> TractResult<Tensor> {
2097    ensure!(
2098        t.shape().last() == Some(&2),
2099        "The last dimension in the tensor shape {:?} must be 2",
2100        t.shape()
2101    );
2102    unsafe {
2103        t.shape.pop();
2104        t.set_datum_type(t.datum_type().complexify()?);
2105        t.update_strides_and_len();
2106        Ok(t)
2107    }
2108}
2109
2110#[cfg(feature = "complex")]
2111pub fn reinterpret_complex_as_inner_dim(mut t: Tensor) -> TractResult<Tensor> {
2112    unsafe {
2113        t.shape.push(2);
2114        t.set_datum_type(t.datum_type().decomplexify()?);
2115        t.update_strides_and_len();
2116        Ok(t)
2117    }
2118}
2119
2120pub fn clip_range_bounds(len: usize, range: impl std::ops::RangeBounds<usize>) -> Range<usize> {
2121    use std::ops::Bound;
2122    let start = match range.start_bound() {
2123        Bound::Included(ix) => *ix,
2124        Bound::Excluded(ix) => ix + 1,
2125        Bound::Unbounded => 0,
2126    };
2127    let end = match range.end_bound() {
2128        Bound::Included(ix) => *ix + 1,
2129        Bound::Excluded(ix) => *ix,
2130        Bound::Unbounded => len,
2131    };
2132    start..end
2133}
2134
2135pub fn natural_strides(shape: &[usize]) -> TVec<isize> {
2136    let mut strides = tvec!();
2137    compute_natural_stride_to(&mut strides, shape);
2138    strides
2139}
2140
2141fn compute_natural_stride_to(strides: &mut TVec<isize>, shape: &[usize]) {
2142    match shape.len() {
2143        0 => (),
2144        1 => strides.push(1),
2145        2 => strides.extend_from_slice(&[shape[1] as isize, 1]),
2146        3 => strides.extend_from_slice(&[(shape[1] * shape[2]) as isize, shape[2] as _, 1]),
2147        4 => strides.extend_from_slice(&[
2148            (shape[1] * shape[2] * shape[3]) as isize,
2149            (shape[2] * shape[3]) as _,
2150            shape[3] as _,
2151            1,
2152        ]),
2153        _ => {
2154            strides.push(1);
2155            for dim in shape.as_ref().iter().skip(1).rev() {
2156                let previous = *strides.last().unwrap();
2157                strides.push(previous * *dim as isize)
2158            }
2159            strides.reverse();
2160        }
2161    }
2162}
2163
2164impl<D: ::ndarray::Dimension, T: Datum> From<Array<T, D>> for Tensor {
2165    fn from(it: Array<T, D>) -> Tensor {
2166        Tensor::from_datum(it.into_dyn())
2167    }
2168}
2169
2170/// Convenient conversion to Tensor.
2171pub trait IntoTensor: Sized {
2172    /// Convert Self to a Tensor.
2173    ///
2174    /// May perform a copy
2175    fn into_tensor(self) -> Tensor;
2176}
2177
2178/// Convenient conversion to Arc<Tensor>.
2179pub trait IntoArcTensor: Sized {
2180    /// Convert Self to a Arc<Tensor>.
2181    ///
2182    /// May perform a copy
2183    fn into_arc_tensor(self) -> Arc<Tensor>;
2184}
2185
2186impl<D: ::ndarray::Dimension, T: Datum> IntoTensor for Array<T, D> {
2187    fn into_tensor(self) -> Tensor {
2188        Tensor::from(self)
2189    }
2190}
2191
2192impl<D: ::ndarray::Dimension, T: Datum> IntoArcTensor for Array<T, D> {
2193    fn into_arc_tensor(self) -> Arc<Tensor> {
2194        Arc::new(Tensor::from(self))
2195    }
2196}
2197
2198impl IntoTensor for Tensor {
2199    fn into_tensor(self) -> Tensor {
2200        self
2201    }
2202}
2203
2204impl IntoTensor for Arc<Tensor> {
2205    fn into_tensor(self) -> Tensor {
2206        Arc::try_unwrap(self).unwrap_or_else(|t| (*t).clone())
2207    }
2208}
2209
2210impl IntoArcTensor for Tensor {
2211    fn into_arc_tensor(self) -> Arc<Tensor> {
2212        Arc::new(self)
2213    }
2214}
2215
2216impl IntoArcTensor for Arc<Tensor> {
2217    fn into_arc_tensor(self) -> Arc<Tensor> {
2218        self
2219    }
2220}
2221
2222#[cfg(test)]
2223mod tests {
2224    use crate::dim::SymbolScope;
2225    use crate::prelude::tensor1;
2226
2227    use super::*;
2228    use litteral::tensor0;
2229    use proptest::collection::vec;
2230    use proptest::prelude::*;
2231
2232    // Regression for sonos/tract#2390: from_raw must reject a content length that
2233    // does not match the declared shape rather than panicking in copy_from_slice.
2234    #[test]
2235    fn from_raw_rejects_length_mismatch() {
2236        // shape [2, 3] of f32 needs 24 bytes; supply 12.
2237        let err = unsafe { Tensor::from_raw_dt(f32::datum_type(), &[2, 3], &[0u8; 12]) }
2238            .expect_err("from_raw must reject a short content buffer, not panic");
2239        assert!(err.to_string().contains("does not match shape"), "unexpected error: {err}");
2240        // Too-long content is rejected as well.
2241        assert!(unsafe { Tensor::from_raw_dt(f32::datum_type(), &[2, 3], &[0u8; 32]) }.is_err());
2242        // Exact match still succeeds.
2243        assert!(unsafe { Tensor::from_raw_dt(f32::datum_type(), &[2, 3], &[0u8; 24]) }.is_ok());
2244    }
2245
2246    #[derive(Debug)]
2247    struct PermuteAxisProblem {
2248        shape: Vec<usize>,
2249        permutation: Vec<usize>,
2250    }
2251
2252    impl Arbitrary for PermuteAxisProblem {
2253        type Strategy = BoxedStrategy<PermuteAxisProblem>;
2254        type Parameters = ();
2255
2256        fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
2257            (0..8usize)
2258                .prop_flat_map(|rank| {
2259                    let permute: Vec<usize> = (0..rank).collect();
2260                    (proptest::collection::vec(1..5usize, rank), Just(permute).prop_shuffle())
2261                })
2262                .prop_map(|(shape, permutation)| PermuteAxisProblem { shape, permutation })
2263                .boxed()
2264        }
2265    }
2266
2267    impl PermuteAxisProblem {
2268        fn input(&self) -> ArrayD<i32> {
2269            let mut i = 0;
2270            ArrayD::from_shape_simple_fn(&*self.shape, || {
2271                i += 1;
2272                i
2273            })
2274            .permuted_axes(&*self.permutation)
2275        }
2276
2277        fn reference(&self) -> Tensor {
2278            let values: Vec<i32> = self.input().iter().copied().collect();
2279            let shape = self.permutation.iter().map(|ix| self.shape[*ix]).collect::<TVec<usize>>();
2280            super::litteral::tensor1(&values).into_shape(&shape).unwrap()
2281        }
2282
2283        fn tract(&self) -> Tensor {
2284            Tensor::from(self.input())
2285        }
2286
2287        fn check(&self) -> proptest::test_runner::TestCaseResult {
2288            prop_assert_eq!(self.tract(), self.reference());
2289            Ok(())
2290        }
2291    }
2292
2293    proptest::proptest! {
2294        #[test]
2295        fn prop(pb: PermuteAxisProblem) {
2296            pb.check().unwrap();
2297        }
2298    }
2299
2300    #[test]
2301    fn t_1_2() {
2302        PermuteAxisProblem { shape: vec![2, 1], permutation: vec![1, 0] }.check().unwrap();
2303    }
2304
2305    #[test]
2306    fn t_2_2() {
2307        PermuteAxisProblem { shape: vec![2, 2], permutation: vec![1, 0] }.check().unwrap();
2308    }
2309
2310    #[derive(Debug)]
2311    struct BroadcastVecToShape {
2312        vec: Vec<f32>,
2313        axis: usize,
2314        shape: TVec<usize>,
2315    }
2316
2317    impl BroadcastVecToShape {
2318        fn check(&self) -> proptest::test_runner::TestCaseResult {
2319            let input = tensor1(&self.vec);
2320            let mut intermediate = tvec![1usize; self.shape.len()];
2321            intermediate[self.axis] = self.vec.len();
2322            let reference = input
2323                .clone()
2324                .into_shape(&intermediate)
2325                .unwrap()
2326                .broadcast_to_shape(&self.shape)
2327                .unwrap();
2328            prop_assert_eq!(
2329                reference,
2330                input.broadcast_vector_to_shape(&self.shape, self.axis).unwrap()
2331            );
2332            Ok(())
2333        }
2334    }
2335
2336    impl Arbitrary for BroadcastVecToShape {
2337        type Strategy = BoxedStrategy<BroadcastVecToShape>;
2338        type Parameters = ();
2339
2340        fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
2341            vec(0usize..5, 0usize..4)
2342                .prop_flat_map(|shape| {
2343                    (vec(-10f32..10f32, 0usize..5), Just(shape.clone()), 0..shape.len() + 1)
2344                })
2345                .prop_map(|(vec, mut shape, axis)| {
2346                    shape.insert(axis, vec.len());
2347                    BroadcastVecToShape { vec, shape: shape.into(), axis }
2348                })
2349                .boxed()
2350        }
2351    }
2352
2353    proptest::proptest! {
2354        #[test]
2355        fn broadcast_vector_to_shape_prop(pb: BroadcastVecToShape) {
2356            pb.check().unwrap()
2357        }
2358    }
2359
2360    #[test]
2361    #[cfg(feature = "complex")]
2362    fn test_reinterpret_inner_dim_as_complex() -> TractResult<()> {
2363        let input = crate::internal::tensor2(&[[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]]);
2364        let cplx_input = reinterpret_inner_dim_as_complex(input)?;
2365        let expected = crate::internal::tensor1(&[
2366            Complex::new(1.0f32, 2.0),
2367            Complex::new(3.0, 4.0),
2368            Complex::new(5.0, 6.0),
2369        ]);
2370        assert_eq!(expected, cplx_input);
2371        Ok(())
2372    }
2373
2374    #[test]
2375    #[cfg(feature = "complex")]
2376    fn test_reinterpret_inner_dim_as_complex_2() -> TractResult<()> {
2377        let input =
2378            crate::internal::tensor3(&[[[1i32, 2], [1, 2]], [[3, 4], [3, 4]], [[5, 6], [5, 6]]]);
2379        let cplx_input = reinterpret_inner_dim_as_complex(input)?;
2380        let expected = crate::internal::tensor2(&[
2381            [Complex::new(1i32, 2), Complex::new(1, 2)],
2382            [Complex::new(3, 4), Complex::new(3, 4)],
2383            [Complex::new(5, 6), Complex::new(5, 6)],
2384        ]);
2385        assert_eq!(expected, cplx_input);
2386        Ok(())
2387    }
2388
2389    #[test]
2390    fn clone_tdim_tensor() {
2391        let symbols = SymbolScope::default();
2392        let a = symbols.sym("a");
2393        let t = tensor0(TDim::from(a));
2394        let _ = t.clone();
2395    }
2396
2397    #[test]
2398    fn ulp_approximation_accepts_within_bound() -> TractResult<()> {
2399        let a = tensor1(&[1.0f32, 2.0, 3.0]);
2400        let b = tensor1(&[
2401            f32::from_bits(1.0f32.to_bits() + 1),
2402            2.0,
2403            f32::from_bits(3.0f32.to_bits() + 2),
2404        ]);
2405        a.close_enough(&b, Approximation::Ulp(2))?;
2406        assert!(a.close_enough(&b, Approximation::Ulp(1)).is_err());
2407        assert_eq!(a.max_ulp_distance(&b)?, (2, Some(2)));
2408        Ok(())
2409    }
2410
2411    #[test]
2412    fn ulp_approximation_uses_the_tensor_own_float_type() -> TractResult<()> {
2413        // One f16 rounding step is ~8192 f32 steps. Measuring in f32 would make an
2414        // adjacent-f16 pair look wildly off, so the comparison must stay in f16.
2415        let one = f16::from_f32(1.0);
2416        let a = tensor1(&[one]);
2417        let b = tensor1(&[f16::from_bits(one.to_bits() + 1)]);
2418        assert_eq!(a.ulp_comparison_dt(), DatumType::F16);
2419        assert_eq!(a.max_ulp_distance(&b)?, (1, Some(0)));
2420        a.close_enough(&b, Approximation::Ulp(1))?;
2421        Ok(())
2422    }
2423
2424    #[test]
2425    fn ulp_approximation_is_scale_free() -> TractResult<()> {
2426        // The same relative error at wildly different magnitudes reads the same,
2427        // which a shared atol cannot do.
2428        let a = tensor1(&[1e-30f32, 1e30]);
2429        let b = tensor1(&[
2430            f32::from_bits(1e-30f32.to_bits() + 1),
2431            f32::from_bits(1e30f32.to_bits() + 1),
2432        ]);
2433        a.close_enough(&b, Approximation::Ulp(1))?;
2434        Ok(())
2435    }
2436
2437    #[test]
2438    fn ulp_approximation_rejects_shape_mismatch() {
2439        let a = tensor1(&[1.0f32, 2.0]);
2440        let b = tensor1(&[1.0f32]);
2441        assert!(a.close_enough(&b, Approximation::Ulp(1000)).is_err());
2442    }
2443
2444    // stack_tensors picks between three copy strategies by block size; they must
2445    // all agree with plain index arithmetic. Sizes 1/2/4/8 cover the typed-copy
2446    // dispatch, and a trailing axis of 1 gives the one-datum blocks that
2447    // FastEnhancer's [1, 256, 1, 2] and DTLN's [1, 2, 128, 2] concats produce.
2448    fn stack_reference<T: Datum + Copy + num_traits::Zero>(
2449        axis: usize,
2450        tensors: &[Tensor],
2451    ) -> Tensor {
2452        let mut shape: TVec<usize> = tensors[0].shape().into();
2453        shape[axis] = tensors.iter().map(|t| t.shape()[axis]).sum();
2454        let mut out = Tensor::zero::<T>(&shape).unwrap();
2455        let outer: usize = shape[..axis].iter().product();
2456        let inner: usize = shape[axis + 1..].iter().product();
2457        let mid = shape[axis];
2458        let ov = unsafe { out.as_slice_mut_unchecked::<T>() };
2459        let mut base = 0;
2460        for t in tensors {
2461            let m = t.shape()[axis];
2462            let tv = unsafe { t.as_slice_unchecked::<T>() };
2463            for o in 0..outer {
2464                for j in 0..m {
2465                    for i in 0..inner {
2466                        ov[(o * mid + base + j) * inner + i] = tv[(o * m + j) * inner + i];
2467                    }
2468                }
2469            }
2470            base += m;
2471        }
2472        out
2473    }
2474
2475    fn ramp<T: Datum + Copy + From<u8>>(shape: &[usize], seed: u8) -> Tensor {
2476        let n: usize = shape.iter().product();
2477        let v: Vec<T> = (0..n).map(|i| T::from(seed.wrapping_add(i as u8))).collect();
2478        Tensor::from_shape(shape, &v).unwrap()
2479    }
2480
2481    macro_rules! stack_agrees_for {
2482        ($name:ident, $t:ty) => {
2483            #[test]
2484            fn $name() {
2485                for shape in [
2486                    tvec!(1usize, 256, 1, 1),
2487                    tvec!(1usize, 2, 128, 1),
2488                    tvec!(1usize, 35, 35, 8),
2489                    tvec!(4usize, 3),
2490                    tvec!(7usize),
2491                ] {
2492                    for axis in 0..shape.len() {
2493                        let a: Tensor = ramp::<$t>(&shape, 1);
2494                        let b: Tensor = ramp::<$t>(&shape, 100);
2495                        let c: Tensor = ramp::<$t>(&shape, 200);
2496                        for n in 1..=3 {
2497                            let ins = [a.clone(), b.clone(), c.clone()][..n].to_vec();
2498                            let got = Tensor::stack_tensors(axis, &ins).unwrap();
2499                            let want = stack_reference::<$t>(axis, &ins);
2500                            assert_eq!(got, want, "shape {shape:?} axis {axis} n {n}");
2501                        }
2502                    }
2503                }
2504            }
2505        };
2506    }
2507
2508    stack_agrees_for!(stack_tensors_agrees_u8, u8);
2509    stack_agrees_for!(stack_tensors_agrees_u16, u16);
2510    stack_agrees_for!(stack_tensors_agrees_u32, u32);
2511    stack_agrees_for!(stack_tensors_agrees_u64, u64);
2512
2513    // A zero extent before the concatenated axis makes `outer` zero; the block
2514    // path divides by it, so it must not be taken.
2515    #[test]
2516    fn stack_tensors_tolerates_a_zero_outer_extent() {
2517        let a = Tensor::zero::<f32>(&[0, 2, 3]).unwrap();
2518        let stacked = Tensor::stack_tensors(2, &[a.clone(), a.clone()]).unwrap();
2519        assert_eq!(stacked.shape(), &[0, 2, 6]);
2520    }
2521
2522    // assign_slice copies one contiguous run per coordinate of the axes before
2523    // `axis`; the runs must agree with plain index arithmetic for every axis,
2524    // including the trailing one where each run is a single datum.
2525    fn assign_slice_reference<T: Datum + Copy>(
2526        dst: &Tensor,
2527        dst_range: Range<usize>,
2528        src: &Tensor,
2529        src_range: Range<usize>,
2530        axis: usize,
2531    ) -> Tensor {
2532        let mut out = dst.clone();
2533        let outer: usize = dst.shape()[..axis].iter().product();
2534        let inner: usize = dst.shape()[axis + 1..].iter().product();
2535        let dst_mid = dst.shape()[axis];
2536        let src_mid = src.shape()[axis];
2537        let sv = unsafe { src.as_slice_unchecked::<T>() };
2538        let ov = unsafe { out.as_slice_mut_unchecked::<T>() };
2539        for o in 0..outer {
2540            for j in 0..dst_range.len() {
2541                for i in 0..inner {
2542                    ov[(o * dst_mid + dst_range.start + j) * inner + i] =
2543                        sv[(o * src_mid + src_range.start + j) * inner + i];
2544                }
2545            }
2546        }
2547        out
2548    }
2549
2550    macro_rules! assign_slice_agrees_for {
2551        ($name:ident, $t:ty) => {
2552            #[test]
2553            fn $name() {
2554                for (shape, axis, dst_mid, src_mid, dst_start, len, src_start) in [
2555                    (tvec!(1usize, 56, 24), 2, 24, 8, 16, 8, 0),
2556                    (tvec!(1usize, 56, 24), 2, 24, 24, 0, 16, 8),
2557                    (tvec!(1usize, 32, 4, 128), 3, 128, 128, 0, 64, 64),
2558                    (tvec!(1usize, 8, 16, 64), 2, 16, 1, 3, 1, 0),
2559                    (tvec!(3usize, 5), 0, 3, 7, 1, 2, 4),
2560                    (tvec!(4usize, 3), 1, 3, 3, 0, 3, 0),
2561                    (tvec!(7usize), 0, 7, 7, 2, 0, 5),
2562                ] {
2563                    let mut dst_shape = shape.clone();
2564                    dst_shape[axis] = dst_mid;
2565                    let mut src_shape = shape.clone();
2566                    src_shape[axis] = src_mid;
2567                    let mut got: Tensor = ramp::<$t>(&dst_shape, 1);
2568                    let src: Tensor = ramp::<$t>(&src_shape, 100);
2569                    let want = assign_slice_reference::<$t>(
2570                        &got,
2571                        dst_start..dst_start + len,
2572                        &src,
2573                        src_start..src_start + len,
2574                        axis,
2575                    );
2576                    got.assign_slice(
2577                        dst_start..dst_start + len,
2578                        &src,
2579                        src_start..src_start + len,
2580                        axis,
2581                    )
2582                    .unwrap();
2583                    assert_eq!(got, want, "shape {dst_shape:?} axis {axis}");
2584                }
2585            }
2586        };
2587    }
2588
2589    assign_slice_agrees_for!(assign_slice_agrees_u8, u8);
2590    assign_slice_agrees_for!(assign_slice_agrees_u16, u16);
2591    assign_slice_agrees_for!(assign_slice_agrees_u32, u32);
2592    assign_slice_agrees_for!(assign_slice_agrees_u64, u64);
2593
2594    // The prefixed assign copies the same runs inside one sub-tensor of each
2595    // side, and the prefixed axes are free to differ in extent.
2596    macro_rules! assign_slice_at_prefix_agrees_for {
2597        ($name:ident, $t:ty) => {
2598            #[test]
2599            fn $name() {
2600                for (shape, prefix, src_lead, src_prefix, axis, start, len, src_start) in [
2601                    (tvec!(3usize, 5, 7), tvec!(2usize), 4, tvec!(3usize), 2, 3, 4, 0),
2602                    (tvec!(3usize, 5, 7), tvec!(0usize), 1, tvec!(0usize), 1, 1, 3, 2),
2603                    (tvec!(2usize, 4, 8, 3), tvec!(1usize, 2), 2, tvec!(0usize, 1), 3, 0, 3, 0),
2604                    (tvec!(4usize, 6), tvec!(), 4, tvec!(), 1, 2, 4, 2),
2605                ] {
2606                    let mut src_shape = shape.clone();
2607                    src_shape[0] = src_lead;
2608                    let mut got: Tensor = ramp::<$t>(&shape, 1);
2609                    let src: Tensor = ramp::<$t>(&src_shape, 100);
2610                    // A prefixed assign is the same assign between the two
2611                    // sub-tensors, which a slice materializes.
2612                    let mut want_sub = sub_tensor(&got, &prefix);
2613                    want_sub
2614                        .assign_slice(
2615                            start..start + len,
2616                            &sub_tensor(&src, &src_prefix),
2617                            src_start..src_start + len,
2618                            axis - prefix.len(),
2619                        )
2620                        .unwrap();
2621                    got.assign_slice_at_prefix(
2622                        &prefix,
2623                        start..start + len,
2624                        &src,
2625                        &src_prefix,
2626                        src_start..src_start + len,
2627                        axis,
2628                    )
2629                    .unwrap();
2630                    assert_eq!(
2631                        sub_tensor(&got, &prefix),
2632                        want_sub,
2633                        "shape {shape:?} prefix {prefix:?} axis {axis}"
2634                    );
2635                }
2636            }
2637        };
2638    }
2639
2640    fn sub_tensor(t: &Tensor, prefix: &[usize]) -> Tensor {
2641        let mut sub = t.clone();
2642        for ix in prefix {
2643            sub = sub.slice(0, *ix, ix + 1).unwrap();
2644            sub.remove_axis(0).unwrap();
2645        }
2646        sub
2647    }
2648
2649    assign_slice_at_prefix_agrees_for!(assign_slice_at_prefix_agrees_u8, u8);
2650    assign_slice_at_prefix_agrees_for!(assign_slice_at_prefix_agrees_u16, u16);
2651    assign_slice_at_prefix_agrees_for!(assign_slice_at_prefix_agrees_u32, u32);
2652    assign_slice_at_prefix_agrees_for!(assign_slice_at_prefix_agrees_u64, u64);
2653
2654    // fill_slice writes one contiguous run per coordinate of the axes between
2655    // the prefix and `axis`; the runs must agree with plain index arithmetic,
2656    // including on the trailing axis where each run is a single datum.
2657    fn fill_slice_reference<T: Datum + Copy>(
2658        data: &Tensor,
2659        prefix: &[usize],
2660        range: Range<usize>,
2661        value: T,
2662        axis: usize,
2663    ) -> Tensor {
2664        let mut out = data.clone();
2665        let shape = data.shape().to_vec();
2666        let inner: usize = shape[axis + 1..].iter().product();
2667        let mid = shape[axis];
2668        let outer: usize = shape[prefix.len()..axis].iter().product();
2669        let at: usize = izip!(prefix, data.strides()).map(|(ix, s)| ix * *s as usize).sum();
2670        let ov = unsafe { out.as_slice_mut_unchecked::<T>() };
2671        for o in 0..outer {
2672            for j in range.clone() {
2673                for i in 0..inner {
2674                    ov[at + (o * mid + j) * inner + i] = value;
2675                }
2676            }
2677        }
2678        out
2679    }
2680
2681    macro_rules! fill_slice_agrees_for {
2682        ($name:ident, $t:ty) => {
2683            #[test]
2684            fn $name() {
2685                for (shape, prefix, axis, start, len) in [
2686                    (tvec!(1usize, 56, 24), tvec!(), 2, 16, 8),
2687                    (tvec!(3usize, 5, 7), tvec!(), 1, 1, 3),
2688                    (tvec!(3usize, 5, 7), tvec!(2usize), 2, 3, 4),
2689                    (tvec!(3usize, 5, 7), tvec!(1usize, 4), 2, 0, 7),
2690                    (tvec!(4usize, 3), tvec!(), 0, 1, 2),
2691                    (tvec!(2usize, 8, 16, 64), tvec!(1usize), 2, 3, 1),
2692                    (tvec!(7usize), tvec!(), 0, 2, 0),
2693                ] {
2694                    let value: $t = 42 as $t;
2695                    let mut got: Tensor = ramp::<$t>(&shape, 1);
2696                    let want =
2697                        fill_slice_reference::<$t>(&got, &prefix, start..start + len, value, axis);
2698                    got.fill_slice_at_prefix(&prefix, start..start + len, &tensor0(value), axis)
2699                        .unwrap();
2700                    assert_eq!(got, want, "shape {shape:?} prefix {prefix:?} axis {axis}");
2701                }
2702            }
2703        };
2704    }
2705
2706    fill_slice_agrees_for!(fill_slice_agrees_u8, u8);
2707    fill_slice_agrees_for!(fill_slice_agrees_u16, u16);
2708    fill_slice_agrees_for!(fill_slice_agrees_u32, u32);
2709    fill_slice_agrees_for!(fill_slice_agrees_u64, u64);
2710
2711    #[test]
2712    fn fill_slice_carries_non_copy_data() {
2713        let strings = |v: [&str; 6]| {
2714            ndarray::Array2::from_shape_vec((2, 3), v.iter().map(|s| s.to_string()).collect())
2715                .unwrap()
2716                .into_tensor()
2717        };
2718        let mut data = strings(["a", "b", "c", "d", "e", "f"]);
2719        data.fill_slice(1..3, &tensor0("x".to_string()), 1).unwrap();
2720        assert_eq!(data, strings(["a", "x", "x", "d", "x", "x"]));
2721    }
2722
2723    #[test]
2724    fn assign_slice_carries_non_copy_data() {
2725        let strings = |v: [&str; 6]| {
2726            ndarray::Array2::from_shape_vec((2, 3), v.iter().map(|s| s.to_string()).collect())
2727                .unwrap()
2728                .into_tensor()
2729        };
2730        let mut dst = strings(["a", "b", "c", "d", "e", "f"]);
2731        let src = ndarray::Array2::from_shape_vec((2, 1), vec!["x".to_string(), "y".to_string()])
2732            .unwrap()
2733            .into_tensor();
2734        dst.assign_slice(1..2, &src, 0..1, 1).unwrap();
2735        assert_eq!(dst, strings(["a", "x", "c", "d", "y", "f"]));
2736    }
2737
2738    // The run-based broadcast must agree with the ndarray view it replaced, over
2739    // leading, middle and trailing broadcast axes and over ranks that grow.
2740    #[test]
2741    fn broadcast_to_shape_agrees_with_the_view() {
2742        for (src, dst) in [
2743            (tvec!(1usize, 8, 1, 7, 4), tvec!(1usize, 8, 4, 7, 4)),
2744            (tvec!(1usize, 1, 1, 7), tvec!(2usize, 3, 5, 7)),
2745            (tvec!(4usize), tvec!(2usize, 3, 5, 4)),
2746            (tvec!(1usize, 5, 3), tvec!(6usize, 5, 3)),
2747            (tvec!(2usize, 3), tvec!(2usize, 3)),
2748            (tvec!(1usize), tvec!(3usize, 1, 2)),
2749            (tvec!(3usize, 1), tvec!(3usize, 0)),
2750        ] {
2751            for dt in [f32::datum_type(), u8::datum_type(), i32::datum_type()] {
2752                let t = Tensor::zero_dt(dt, &src).unwrap().cast_to_dt(dt).unwrap().into_owned();
2753                let got = t.broadcast_to_shape(&dst).unwrap();
2754                let want = dispatch_datum!(Tensor::broadcast_to_shape_t(dt)(&t, &dst)).unwrap();
2755                assert_eq!(got.shape(), &*dst, "{src:?} -> {dst:?}");
2756                assert_eq!(got, want, "{src:?} -> {dst:?} {dt:?}");
2757            }
2758        }
2759    }
2760
2761    #[test]
2762    fn broadcast_to_shape_carries_values_and_rejects_mismatches() {
2763        let t = tensor2(&[[1u8, 2, 3], [4, 5, 6]]);
2764        let got = t.clone().into_shape(&[2, 1, 3]).unwrap().broadcast_to_shape(&[2, 2, 3]).unwrap();
2765        assert_eq!(got, tensor3(&[[[1u8, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6]]]));
2766        assert!(t.broadcast_to_shape(&[3, 3]).is_err());
2767        assert!(t.broadcast_to_shape(&[3]).is_err());
2768    }
2769
2770    #[test]
2771    fn slice_keeps_the_datum_type_and_the_values() {
2772        let t = ramp::<u32>(&tvec!(2usize, 3, 4), 0);
2773        let got = t.slice(1, 1, 3).unwrap();
2774        let mut want = Tensor::zero::<u32>(&[2, 2, 4]).unwrap();
2775        want.assign_slice(0..2, &t, 1..3, 1).unwrap();
2776        assert_eq!(got, want);
2777        let quantized = Tensor::zero_dt(
2778            i8::datum_type().quantize(QParams::ZpScale { zero_point: 3, scale: 0.5 }),
2779            &[2, 4],
2780        )
2781        .unwrap();
2782        assert_eq!(quantized.slice(1, 0, 2).unwrap().datum_type(), quantized.datum_type());
2783    }
2784
2785    #[test]
2786    fn ulp_bounds_are_distinguished_by_equality() {
2787        assert_eq!(Approximation::Ulp(1), Approximation::Ulp(1));
2788        assert_ne!(Approximation::Ulp(1), Approximation::Ulp(2));
2789    }
2790}