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