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, Zero};
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
188impl Tensor {
189    #[inline]
190    fn plain_storage(&self) -> &PlainStorage {
191        self.storage.as_plain().expect("Non-plain storage")
192    }
193
194    #[inline]
195    fn plain_storage_mut(&mut self) -> &mut PlainStorage {
196        self.storage.as_plain_mut().expect("Non-plain storage")
197    }
198
199    pub fn storage_as<T: TensorStorage>(&self) -> Option<&T> {
200        self.storage.as_storage().downcast_ref::<T>()
201    }
202
203    pub fn try_storage_as<T: TensorStorage>(&self) -> TractResult<&T> {
204        self.storage_as::<T>().context("Unexpected tensor storage type")
205    }
206
207    pub fn from_storage(
208        dt: DatumType,
209        shape: &[usize],
210        storage: impl TensorStorage + 'static,
211    ) -> Tensor {
212        let len = shape.iter().product::<usize>();
213        let strides = Self::natural_strides(shape);
214        Tensor {
215            dt,
216            shape: shape.into(),
217            strides,
218            len,
219            storage: StorageKind::Exotic(Box::new(storage)),
220        }
221    }
222
223    /// Returns an immutable [`PlainView`] if this tensor has plain storage.
224    #[inline]
225    pub fn as_plain(&self) -> Option<PlainView<'_>> {
226        let storage = self.storage.as_plain()?;
227        Some(PlainView::new(self, storage))
228    }
229
230    /// Returns an immutable [`PlainView`], or an error if storage is not plain.
231    #[inline]
232    pub fn try_as_plain(&self) -> TractResult<PlainView<'_>> {
233        self.as_plain().context("Tensor storage is not plain")
234    }
235
236    /// Returns `true` if this tensor uses plain (contiguous) storage.
237    #[inline]
238    pub fn is_plain(&self) -> bool {
239        self.storage.as_plain().is_some()
240    }
241
242    /// Returns `true` if this tensor uses exotic (non-plain) storage.
243    #[inline]
244    pub fn is_exotic(&self) -> bool {
245        !self.is_plain()
246    }
247
248    /// Build the `ExoticFact` matching this tensor's storage, or `None` for plain tensors.
249    pub fn exotic_fact(&self) -> TractResult<Option<Box<dyn crate::exotic::ExoticFact>>> {
250        self.storage.as_storage().exotic_fact(&self.shape)
251    }
252
253    /// Returns a mutable [`PlainViewMut`] if this tensor has plain storage.
254    #[inline]
255    pub fn as_plain_mut(&mut self) -> Option<PlainViewMut<'_>> {
256        let storage = self.storage.as_plain_mut()?;
257        Some(PlainViewMut::new(self.dt, &self.shape, &self.strides, self.len, storage))
258    }
259
260    /// Returns a mutable [`PlainViewMut`], or an error if storage is not plain.
261    #[inline]
262    pub fn try_as_plain_mut(&mut self) -> TractResult<PlainViewMut<'_>> {
263        self.as_plain_mut().context("Tensor storage is not plain")
264    }
265
266    /// Create an uninitialized tensor (dt as type paramater).
267    #[inline]
268    pub unsafe fn uninitialized<T: Datum>(shape: &[usize]) -> TractResult<Tensor> {
269        unsafe { Self::uninitialized_dt(T::datum_type(), shape) }
270    }
271
272    /// Create an uninitialized tensor (dt as regular parameter).
273    #[inline]
274    pub unsafe fn uninitialized_dt(dt: DatumType, shape: &[usize]) -> TractResult<Tensor> {
275        unsafe { Self::uninitialized_aligned_dt(dt, shape, vector_size()) }
276    }
277
278    /// Create an uninitialized tensor with a given alignment (in bytes).
279    #[inline]
280    pub unsafe fn uninitialized_aligned<T: Datum>(
281        shape: &[usize],
282        alignment: usize,
283    ) -> TractResult<Tensor> {
284        unsafe { Self::uninitialized_aligned_dt(T::datum_type(), shape, alignment) }
285    }
286
287    /// Create an uninitialized tensor with a given alignment (in bytes).
288    pub unsafe fn uninitialized_aligned_dt(
289        dt: DatumType,
290        shape: &[usize],
291        alignment: usize,
292    ) -> TractResult<Tensor> {
293        let bytes = shape.iter().cloned().product::<usize>() * dt.size_of();
294        let storage = StorageKind::Plain(PlainStorage::from(unsafe {
295            Blob::new_for_size_and_align(bytes, alignment)
296        }));
297        let mut tensor = Tensor { strides: tvec!(), dt, shape: shape.into(), storage, len: 0 };
298        if tensor.shape.len() == 0 {
299            tensor.len = 1;
300        } else {
301            tensor.update_strides_and_len();
302        }
303        if !tensor.storage.is_empty() {
304            if dt == String::datum_type() || dt == Blob::datum_type() {
305                // assumes zero-initialized string and blob are valid
306                tensor.plain_storage_mut().as_bytes_mut().fill(0);
307            } else if dt == TDim::datum_type() {
308                unsafe {
309                    tensor
310                        .as_slice_mut_unchecked::<TDim>()
311                        .iter_mut()
312                        .for_each(|dim| std::ptr::write(dim, TDim::zero()))
313                }
314            } else if cfg!(debug_assertions) {
315                assert!(dt.is_copy());
316                if dt == DatumType::F32 {
317                    tensor.fill_t(f32::NAN).unwrap();
318                } else {
319                    // safe, non copy types have been dealt with
320                    tensor.as_bytes_mut().iter_mut().for_each(|x| *x = (-1i8) as u8);
321                }
322            }
323        }
324        Ok(tensor)
325    }
326
327    pub fn stack_tensors(
328        axis: usize,
329        tensors: &[impl std::borrow::Borrow<Tensor>],
330    ) -> TractResult<Tensor> {
331        ensure!(tensors.len() > 0);
332        let rank = tensors[0].borrow().rank();
333        ensure!(axis < rank);
334        ensure!(tensors.iter().all(|t| t.borrow().rank() == rank));
335        let dt = tensors[0].borrow().datum_type();
336        ensure!(tensors.iter().all(|t| t.borrow().datum_type() == dt));
337        let mut shape: TVec<usize> = tensors[0].borrow().shape().into();
338        for ax in 0..rank {
339            if ax != axis {
340                ensure!(tensors.iter().all(|t| t.borrow().shape()[ax] == shape[ax]));
341            }
342        }
343        shape[axis] = tensors.iter().map(|v| v.borrow().shape()[axis]).sum();
344        unsafe {
345            let mut result = Tensor::uninitialized_dt(dt, &shape)?;
346            // Every input keeps the same trailing block, so one outer stride walks
347            // them alongside the result and each contribution stays contiguous.
348            let outer: usize = shape[..axis].iter().product();
349            let out_stride = shape[axis..].iter().product::<usize>() * dt.size_of();
350            if dt.is_copy() && tensors.iter().all(|t| t.borrow().storage.as_plain().is_some()) {
351                let out = result.plain_storage_mut().as_mut_ptr();
352                let mut offset = 0isize;
353                for v in tensors {
354                    let v = v.borrow();
355                    let block = v.storage.byte_len() / outer;
356                    let src = v.plain_storage().as_ptr();
357                    for o in 0..outer {
358                        std::ptr::copy_nonoverlapping(
359                            src.add(o * block),
360                            out.offset(offset + (o * out_stride) as isize),
361                            block,
362                        );
363                    }
364                    offset += block as isize;
365                }
366            } else {
367                let mut offset = 0;
368                for t in tensors {
369                    let t = t.borrow();
370                    let len = t.shape()[axis];
371                    result.assign_slice_from_resolved(offset..offset + len, t, 0..len, axis);
372                    offset += len;
373                }
374            }
375
376            Ok(result)
377        }
378    }
379
380    pub fn clear<T: Datum + num_traits::Zero + Clone>(&mut self) -> TractResult<()> {
381        self.fill_t(T::zero())
382    }
383
384    pub fn zero<T: Datum + num_traits::Zero>(shape: &[usize]) -> TractResult<Tensor> {
385        unsafe {
386            let mut t = Tensor::uninitialized::<T>(shape)?;
387            t.clear::<T>()?;
388            Ok(t)
389        }
390    }
391
392    pub fn zero_scalar<T: Datum + num_traits::Zero>() -> TractResult<Tensor> {
393        Tensor::zero::<T>(&[])
394    }
395
396    pub fn zero_scalar_dt(dt: DatumType) -> TractResult<Tensor> {
397        Tensor::zero_dt(dt, &[])
398    }
399
400    pub fn zero_dt(dt: DatumType, shape: &[usize]) -> TractResult<Tensor> {
401        Tensor::zero_aligned_dt(dt, shape, vector_size())
402    }
403
404    pub fn fill_t<T: Datum + Clone>(&mut self, value: T) -> TractResult<()> {
405        self.try_as_plain_mut()?
406            .as_slice_mut::<T>()?
407            .iter_mut()
408            .for_each(|item| *item = value.clone());
409        Ok(())
410    }
411
412    pub fn zero_aligned_dt(
413        dt: DatumType,
414        shape: &[usize],
415        alignment: usize,
416    ) -> TractResult<Tensor> {
417        if shape.iter().product::<usize>() == 0 {
418            unsafe { return Tensor::uninitialized_dt(dt, shape) };
419        }
420        if dt.is_quantized() {
421            unsafe {
422                let mut t = Tensor::uninitialized_dt(dt, shape)?;
423                let zp = dt.zp_scale().0;
424                match dt.unquantized() {
425                    DatumType::I8 => t
426                        .try_as_plain_mut()?
427                        .as_slice_mut::<i8>()?
428                        .iter_mut()
429                        .for_each(|item| *item = zp as _),
430                    DatumType::U8 => t
431                        .try_as_plain_mut()?
432                        .as_slice_mut::<u8>()?
433                        .iter_mut()
434                        .for_each(|item| *item = zp as _),
435                    DatumType::I32 => t
436                        .try_as_plain_mut()?
437                        .as_slice_mut::<i32>()?
438                        .iter_mut()
439                        .for_each(|item| *item = zp as _),
440                    _ => unreachable!(),
441                }
442                Ok(t)
443            }
444        } else if dt == DatumType::Bool {
445            let mut t = unsafe { Tensor::uninitialized_dt(dt, shape)? };
446            t.fill_t::<bool>(false)?;
447            Ok(t)
448        } else {
449            dispatch_zerolike!(Self::zero_aligned(dt)(shape, alignment))
450        }
451    }
452
453    pub fn zero_aligned<T: Datum + num_traits::Zero>(
454        shape: &[usize],
455        alignment: usize,
456    ) -> TractResult<Tensor> {
457        unsafe {
458            let mut tensor = Self::uninitialized_aligned::<T>(shape, alignment)?;
459            tensor.clear::<T>()?;
460            Ok(tensor)
461        }
462    }
463
464    /// Create a tensor with a given shape and a slice of elements.
465    /// The data is copied and aligned to size of T.
466    pub fn from_shape<T: Datum + Copy>(shape: &[usize], data: &[T]) -> TractResult<Tensor> {
467        Self::from_shape_align(shape, data, vector_size())
468    }
469
470    /// Create a tensor with a given shape and a slice of elements.
471    /// The data is copied and aligned to given alignment.
472    pub fn from_shape_align<T: Datum + Copy>(
473        shape: &[usize],
474        data: &[T],
475        align: usize,
476    ) -> TractResult<Tensor> {
477        ensure!(
478            data.len() == shape.iter().product::<usize>(),
479            "Shape product must be equal to data length"
480        );
481        unsafe {
482            let bytes = std::slice::from_raw_parts(
483                data.as_ptr() as *const u8,
484                data.len() * T::datum_type().size_of(),
485            );
486            let dt = T::datum_type();
487            Self::from_raw_dt_align(dt, shape, bytes, align)
488        }
489    }
490
491    /// Create a tensor from raw data.
492    ///
493    /// It copies the data, aligning it to the size of T.
494    pub unsafe fn from_raw<T: Datum>(shape: &[usize], content: &[u8]) -> TractResult<Tensor> {
495        unsafe { Tensor::from_raw_dt(T::datum_type(), shape, content) }
496    }
497
498    pub unsafe fn from_raw_aligned<T: Datum>(
499        shape: &[usize],
500        content: &[u8],
501        align: usize,
502    ) -> TractResult<Tensor> {
503        unsafe { Tensor::from_raw_dt_align(T::datum_type(), shape, content, align) }
504    }
505
506    pub unsafe fn from_raw_dt(
507        dt: DatumType,
508        shape: &[usize],
509        content: &[u8],
510    ) -> TractResult<Tensor> {
511        unsafe { Self::from_raw_dt_align(dt, shape, content, vector_size()) }
512    }
513
514    pub unsafe fn from_raw_dt_align(
515        dt: DatumType,
516        shape: &[usize],
517        content: &[u8],
518        align: usize,
519    ) -> TractResult<Tensor> {
520        let mut tensor = unsafe { Tensor::uninitialized_aligned_dt(dt, shape, align) }?;
521        let expected = tensor.as_bytes().len();
522        ensure!(
523            content.len() == expected,
524            "Raw tensor data length ({}) does not match shape {:?} of {:?} ({} bytes)",
525            content.len(),
526            shape,
527            dt,
528            expected
529        );
530        tensor.as_bytes_mut().copy_from_slice(content);
531        Ok(tensor)
532    }
533
534    pub unsafe fn from_slice_align<T: Datum>(content: &[T], align: usize) -> TractResult<Tensor> {
535        let bytes = if content.len() == 0 {
536            &[]
537        } else {
538            unsafe {
539                std::slice::from_raw_parts(
540                    content.as_ptr() as *const u8,
541                    content.len() * T::datum_type().size_of(),
542                )
543            }
544        };
545        unsafe { Self::from_raw_dt_align(T::datum_type(), &[content.len()], bytes, align) }
546    }
547
548    /// Get the number of dimensions (or axes) of the tensor.
549    #[inline]
550    pub fn rank(&self) -> usize {
551        self.shape.len()
552    }
553
554    /// Get the shape of the tensor.
555    #[inline]
556    pub fn shape(&self) -> &[usize] {
557        &self.shape
558    }
559
560    /// Get the number of values in the tensor.
561    #[inline]
562    #[allow(clippy::len_without_is_empty)]
563    pub fn len(&self) -> usize {
564        self.len
565    }
566
567    /// Get the number of valeus in the tensor.
568    #[inline]
569    #[allow(clippy::len_without_is_empty)]
570    pub fn volume(&self) -> usize {
571        self.len
572    }
573
574    /// Get the shape of the tensor.
575    #[inline]
576    pub fn strides(&self) -> &[isize] {
577        &self.strides
578    }
579
580    fn update_strides_and_len(&mut self) {
581        self.strides.clear();
582        if self.shape.len() == 0 {
583            self.len = 1;
584            return;
585        }
586        compute_natural_stride_to(&mut self.strides, &self.shape);
587        self.len = unsafe { *self.strides.get_unchecked(0) as usize * self.shape.get_unchecked(0) };
588    }
589
590    /// Force the tensor shape, no consistency check.
591    pub unsafe fn set_shape_unchecked(&mut self, shape: &[usize]) {
592        if shape != &*self.shape {
593            self.shape.clear();
594            self.shape.extend_from_slice(shape);
595            self.update_strides_and_len();
596        }
597    }
598
599    /// Force the tensor shape and strides, no consistency check.
600    pub unsafe fn set_geometry_unchecked(&mut self, shape: &[usize], strides: &[isize]) {
601        self.shape.clear();
602        self.shape.extend_from_slice(shape);
603        self.strides.clear();
604        self.strides.extend_from_slice(strides);
605    }
606
607    /// Force the tensor shape.
608    pub fn set_shape(&mut self, shape: &[usize]) -> TractResult<()> {
609        if self.len() != shape.iter().product::<usize>() {
610            bail!("Invalid reshape {:?} to {:?}", self.shape, shape);
611        }
612        unsafe { self.set_shape_unchecked(shape) }
613        Ok(())
614    }
615
616    pub fn permute_axes(self, axes: &[usize]) -> TractResult<Tensor> {
617        ensure!(axes.iter().duplicates().next().is_none());
618        ensure!(axes.iter().all(|a| *a < self.rank()));
619        unsafe {
620            #[inline]
621            unsafe fn permute<T: Datum>(axes: &[usize], input: Tensor) -> Tensor {
622                unsafe { input.into_array_unchecked::<T>().permuted_axes(axes).into_tensor() }
623            }
624            let dt = self.datum_type();
625            let mut t = dispatch_datum_by_size!(permute(self.datum_type())(axes, self));
626            t.set_datum_type(dt);
627            Ok(t)
628        }
629    }
630
631    pub fn move_axis(self, from: usize, to: usize) -> TractResult<Tensor> {
632        let mut permutation: Vec<usize> = (0..self.rank()).collect();
633        permutation.remove(from);
634        permutation.insert(to, from);
635        self.permute_axes(&permutation)
636    }
637
638    pub fn collapse_axis_with_next(mut self, axis: usize) -> Tensor {
639        let removed = self.shape.remove(axis + 1);
640        self.shape[axis] *= removed;
641        self.update_strides_and_len();
642        self
643    }
644
645    pub fn split_axis(mut self, axis: usize, outer_dim: usize) -> TractResult<Tensor> {
646        if !self.shape[axis].is_multiple_of(outer_dim) {
647            bail!(
648                "Invalid axis split, shape is {:?}, axis split at {}, outer {}",
649                self.shape,
650                axis,
651                outer_dim
652            );
653        }
654        self.shape.insert(axis + 1, self.shape[axis] / outer_dim);
655        self.shape[axis] = outer_dim;
656        self.update_strides_and_len();
657        Ok(self)
658    }
659
660    /// Reshape the tensor to `shape`.
661    pub fn into_shape(mut self, shape: &[usize]) -> TractResult<Tensor> {
662        self.set_shape(shape)?;
663        Ok(self)
664    }
665
666    pub fn insert_axis(&mut self, axis: usize) -> TractResult<()> {
667        self.shape.insert(axis, 1);
668        self.strides.insert(axis, self.strides.get(axis).copied().unwrap_or(1));
669        Ok(())
670    }
671
672    pub fn remove_axis(&mut self, axis: usize) -> TractResult<()> {
673        ensure!(self.shape[axis] == 1, "Remove a non-1 axis: axis {} in {:?}", axis, self);
674        self.shape.remove(axis);
675        self.strides.remove(axis);
676        Ok(())
677    }
678
679    pub fn broadcast_into_rank(mut self, rank: usize) -> TractResult<Tensor> {
680        self.broadcast_to_rank(rank)?;
681        self.update_strides_and_len();
682        Ok(self)
683    }
684
685    pub fn broadcast_to_rank(&mut self, rank: usize) -> TractResult<()> {
686        if rank < self.rank() {
687            bail!("Can only broadcast to higher rank")
688        }
689        while self.shape.len() < rank {
690            self.shape.insert(0, 1)
691        }
692        self.update_strides_and_len();
693        Ok(())
694    }
695
696    pub fn broadcast_scalar_to_shape(&self, shape: &[usize]) -> TractResult<Tensor> {
697        if self.rank() > 0 {
698            bail!("broadcast_scalar_to_shape called on {:?}, which is not a salar", self);
699        }
700        unsafe fn make<T: Datum>(src: &Tensor, dst: &mut Tensor) {
701            unsafe {
702                let value: &T = src.to_scalar_unchecked::<T>();
703                dst.as_slice_mut_unchecked::<T>().iter_mut().for_each(|item| *item = value.clone())
704            };
705        }
706        unsafe {
707            let mut t = Tensor::uninitialized_dt(self.datum_type(), shape)?;
708            dispatch_datum_by_size!(make(self.datum_type())(self, &mut t));
709            Ok(t)
710        }
711    }
712
713    fn broadcast_to_shape_t<T: Datum>(&self, shape: &[usize]) -> TractResult<Tensor> {
714        unsafe {
715            let view = self.to_array_view_unchecked::<T>();
716            let mut output = view
717                .broadcast(shape)
718                .with_context(|| format!("Broadcasting {view:?} to {shape:?}"))?
719                .into_owned()
720                .into_tensor();
721            output.set_datum_type(self.datum_type());
722            Ok(output)
723        }
724    }
725
726    pub fn broadcast_to_shape(&self, shape: &[usize]) -> TractResult<Tensor> {
727        dispatch_datum!(Self::broadcast_to_shape_t(self.dt)(self, shape))
728    }
729
730    pub fn broadcast_vector_to_shape(&self, shape: &[usize], axis: usize) -> TractResult<Tensor> {
731        ensure!(self.rank() == 1);
732        ensure!(shape[axis] == self.len());
733        if !self.datum_type().is_copy() {
734            let mut vec_shape = vec![1; shape.len()];
735            vec_shape[axis] = self.len();
736            return self.clone().into_shape(&vec_shape)?.broadcast_to_shape(shape);
737        }
738        unsafe {
739            let mut output = Tensor::uninitialized_dt(self.datum_type(), shape)?;
740            if output.len() == 0 {
741                return Ok(output);
742            }
743            let inner_len = shape[axis + 1..].iter().product::<usize>();
744
745            unsafe fn splat<T>(input: &Tensor, output: &mut Tensor, inner_len: usize)
746            where
747                T: Datum + Copy,
748            {
749                unsafe {
750                    for ix in 0..input.len() {
751                        let value: T = input.as_slice_unchecked()[ix];
752                        output.as_slice_mut_unchecked::<T>()[ix * inner_len..(ix + 1) * inner_len]
753                            .iter_mut()
754                            .for_each(|item| *item = value);
755                    }
756                }
757            }
758            dispatch_copy_by_size!(splat(self.datum_type())(&self, &mut output, inner_len));
759
760            let outer_len = shape[0..axis].iter().product::<usize>();
761            let repeat_bytes_len = inner_len * self.as_bytes().len();
762            let bytes = output.as_bytes_mut();
763            for ix in 1..outer_len {
764                bytes.copy_within(0..repeat_bytes_len, ix * repeat_bytes_len);
765            }
766
767            Ok(output)
768        }
769    }
770
771    pub fn assign_slice(
772        &mut self,
773        range: impl std::ops::RangeBounds<usize>,
774        src: &Tensor,
775        src_range: impl std::ops::RangeBounds<usize>,
776        axis: usize,
777    ) -> TractResult<()> {
778        ensure!(self.rank() == src.rank());
779        ensure!(axis < self.rank());
780        let range = clip_range_bounds(self.shape[axis], range);
781        let src_range = clip_range_bounds(src.shape[axis], src_range);
782        ensure!(
783            src.datum_type() == self.datum_type(),
784            "Attempt to assign into {:?} from {:?}, datum type mismatch",
785            self.datum_type(),
786            src.datum_type()
787        );
788        ensure!(
789            src_range.len() == range.len(),
790            "Attempt to assign a range of {:?} from a range of {:?}",
791            range,
792            src_range,
793        );
794        ensure!(
795            itertools::izip!(0.., self.shape(), src.shape())
796                .all(|(ix, dst, src)| ix == axis || src == dst),
797            "Attempt to assign a {}-axis range of {:?} from a range of {:?}",
798            axis,
799            self,
800            src
801        );
802        ensure!(
803            src_range.end <= src.shape()[axis],
804            "Assigning from invalid slice (axis {}, {:?}) of {:?}",
805            axis,
806            src_range,
807            src
808        );
809        ensure!(
810            range.end <= self.shape()[axis],
811            "Assigning to invalid slice (axis {}, {:?}) of {:?}",
812            axis,
813            range,
814            self
815        );
816        unsafe { self.assign_slice_from_resolved(range, src, src_range, axis) };
817        Ok(())
818    }
819
820    pub unsafe fn assign_slice_unchecked(
821        &mut self,
822        range: impl std::ops::RangeBounds<usize>,
823        src: &Tensor,
824        src_range: impl std::ops::RangeBounds<usize>,
825        axis: usize,
826    ) {
827        let range = clip_range_bounds(self.shape[axis], range);
828        let src_range = clip_range_bounds(src.shape[axis], src_range);
829        unsafe { self.assign_slice_from_resolved(range, src, src_range, axis) };
830    }
831
832    #[allow(clippy::ptr_eq)]
833    unsafe fn assign_slice_from_resolved(
834        &mut self,
835        range: std::ops::Range<usize>,
836        src: &Tensor,
837        src_range: std::ops::Range<usize>,
838        axis: usize,
839    ) {
840        unsafe {
841            use ndarray::Slice;
842            unsafe fn assign_slice_t<T: Datum>(
843                to: &mut Tensor,
844                to_range: Range<usize>,
845                from: &Tensor,
846                from_range: Range<usize>,
847                axis: usize,
848            ) {
849                unsafe {
850                    to.to_array_view_mut_unchecked::<T>()
851                        .slice_axis_mut(Axis(axis), Slice::from(to_range))
852                        .assign(
853                            &from
854                                .to_array_view_unchecked::<T>()
855                                .slice_axis(Axis(axis), Slice::from(from_range)),
856                        )
857                }
858            }
859            if self.datum_type().is_copy() && self.shape[..axis].iter().all(|d| *d == 1) {
860                let stride = self.strides[axis] as usize * self.datum_type().size_of();
861                let dst_start = (stride * range.start) as isize;
862                let src_start = (stride * src_range.start) as isize;
863                let len = stride * range.len();
864                if len > 0 {
865                    if self.plain_storage().as_ptr() != src.plain_storage().as_ptr() {
866                        std::ptr::copy_nonoverlapping(
867                            src.plain_storage().as_ptr().offset(src_start),
868                            self.plain_storage_mut().as_mut_ptr().offset(dst_start),
869                            len,
870                        );
871                    } else {
872                        std::ptr::copy(
873                            src.plain_storage().as_ptr().offset(src_start),
874                            self.plain_storage_mut().as_mut_ptr().offset(dst_start),
875                            len,
876                        );
877                    }
878                }
879            } else {
880                dispatch_datum!(assign_slice_t(self.datum_type())(
881                    self, range, src, src_range, axis
882                ));
883            }
884        }
885    }
886
887    /// Get the datum type of the tensor.
888    #[inline]
889    pub fn datum_type(&self) -> DatumType {
890        self.dt
891    }
892
893    /// Set the datum type of the tensor.
894    #[inline]
895    pub unsafe fn set_datum_type(&mut self, dt: DatumType) {
896        self.dt = dt
897    }
898
899    /// Dump the tensor in a human readable form.
900    ///
901    /// `force_full` will force the tensor to be dump in full even if it is big.
902    pub fn dump(&self, force_full: bool) -> TractResult<String> {
903        if self.is_exotic() {
904            return Ok(format!(
905                "{},{:?} (non-plain storage)",
906                self.shape.iter().join(","),
907                self.dt,
908            ));
909        }
910        unsafe fn dump_t<D: Datum>(tensor: &Tensor, n: usize) -> String {
911            unsafe {
912                if let Some(qp) = tensor.datum_type().qparams() {
913                    let integers = tensor.cast_to::<i32>().unwrap();
914                    integers.as_slice_unchecked::<i32>()[0..n]
915                        .iter()
916                        .map(|x| format!("[{}]({})", x, qp.dq(*x)))
917                        .join(", ")
918                } else {
919                    tensor.as_slice_unchecked::<D>()[0..n].iter().join(", ")
920                }
921            }
922        }
923        unsafe {
924            let trunc = self.len() > 12 && !force_full;
925            let data = dispatch_datum!(dump_t(self.datum_type())(
926                self,
927                if trunc { 12 } else { self.len() }
928            ));
929            Ok(format!(
930                "{},{:?} {}{}",
931                self.shape.iter().join(","),
932                self.dt,
933                data,
934                if trunc { "..." } else { "" }
935            ))
936        }
937    }
938
939    /// Compare two tensors, allowing for rounding errors.
940    pub fn close_enough(
941        &self,
942        other: &Self,
943        approx: impl Into<Approximation> + std::fmt::Debug,
944    ) -> TractResult<()> {
945        let approx = approx.into();
946        if self.shape() != other.shape() {
947            bail!("Shape mismatch {:?} != {:?}", self.shape(), other.shape())
948        }
949        if let Approximation::Ulp(max_ulp) = approx {
950            return self.ulp_close_enough(other, max_ulp);
951        }
952        let (atol, rtol, outliers) = approx.atol_rtol_outliers(&self.datum_type());
953        let ma = self.cast_to::<f32>()?;
954        let ma = ma.to_plain_array_view::<f32>()?;
955        let mb = other.cast_to::<f32>()?;
956        let mb = mb.to_plain_array_view::<f32>()?;
957        let mut first_outlier = None;
958        let mut outliers_count = 0;
959        ndarray::indices_of(&ma).into_iter().for_each(|indices| {
960            let a = ma[&indices];
961            let b = mb[&indices];
962            if !((a.is_nan() && b.is_nan())
963                || (a.is_infinite() && b.is_infinite() && a.signum() == b.signum())
964                || (a - b).abs() <= atol as f32 + rtol as f32 * b.abs())
965            {
966                if outliers_count == 0 {
967                    first_outlier = Some(indices.as_array_view().to_vec());
968                }
969                outliers_count += 1;
970            }
971        });
972        if self.volume() > 0 && outliers_count as f64 / self.volume() as f64 > outliers {
973            let indices = first_outlier.unwrap();
974            let a = ma[&*indices];
975            let b = mb[&*indices];
976            let ulp = self
977                .max_ulp_distance(other)
978                .map(|(d, _)| format!("{d}"))
979                .unwrap_or_else(|_| "n/a".to_string());
980            bail!(
981                "Mismatch. First outlier: {:?} for {:?}) at {:?} {} != {}. Outliers: {} / {} = {:0.5} > {:0.5}. Max ULP ({:?}): {}.",
982                approx,
983                self.datum_type(),
984                indices,
985                a,
986                b,
987                outliers_count,
988                self.volume(),
989                outliers_count as f64 / self.volume() as f64,
990                outliers,
991                self.ulp_comparison_dt(),
992                ulp,
993            );
994        }
995        Ok(())
996    }
997
998    /// The float type ULP distances against this tensor are measured in.
999    ///
1000    /// Float tensors are compared in their own type, so an f16 comparison stays an
1001    /// f16 comparison. Anything else falls back to f32, matching what
1002    /// `close_enough` does for its tolerance check.
1003    pub fn ulp_comparison_dt(&self) -> DatumType {
1004        match self.datum_type() {
1005            dt @ (DatumType::F16 | DatumType::F32 | DatumType::F64) => dt,
1006            _ => DatumType::F32,
1007        }
1008    }
1009
1010    /// Largest integer ULP distance between `self` and `other`, and the flat index
1011    /// where it occurs.
1012    ///
1013    /// Comparison happens in [`Self::ulp_comparison_dt`]. See [`crate::ulp`] for
1014    /// the exact convention around signed zeros, infinities and NaN.
1015    pub fn max_ulp_distance(&self, other: &Self) -> TractResult<(u64, Option<usize>)> {
1016        if self.shape() != other.shape() {
1017            bail!("Shape mismatch {:?} != {:?}", self.shape(), other.shape())
1018        }
1019        let dt = self.ulp_comparison_dt();
1020        let a = self.cast_to_dt(dt)?;
1021        let b = other.cast_to_dt(dt)?;
1022        fn worst<D: Datum + crate::ulp::UlpFloat>(
1023            a: &Tensor,
1024            b: &Tensor,
1025        ) -> TractResult<(u64, Option<usize>)> {
1026            let a = a.to_plain_array_view::<D>()?;
1027            let b = b.to_plain_array_view::<D>()?;
1028            Ok(crate::ulp::max_ulp_distance(a.iter().copied(), b.iter().copied()))
1029        }
1030        match dt {
1031            DatumType::F16 => worst::<f16>(&a, &b),
1032            DatumType::F32 => worst::<f32>(&a, &b),
1033            DatumType::F64 => worst::<f64>(&a, &b),
1034            dt => bail!("No ULP comparison for {dt:?}"),
1035        }
1036    }
1037
1038    /// Compare two tensors by integer ULP distance, accepting a distance up to
1039    /// `max_ulp`.
1040    fn ulp_close_enough(&self, other: &Self, max_ulp: u64) -> TractResult<()> {
1041        let (worst, at) = self.max_ulp_distance(other)?;
1042        if worst <= max_ulp {
1043            return Ok(());
1044        }
1045        let dt = self.ulp_comparison_dt();
1046        let indices = at
1047            .map(|flat| {
1048                let mut rest = flat;
1049                let mut indices = vec![0; self.rank()];
1050                for (ix, dim) in self.shape().iter().enumerate().rev() {
1051                    indices[ix] = rest % dim;
1052                    rest /= dim;
1053                }
1054                format!("{indices:?}")
1055            })
1056            .unwrap_or_else(|| "?".to_string());
1057        let a = self.cast_to::<f64>()?;
1058        let b = other.cast_to::<f64>()?;
1059        let (a, b) = (a.to_plain_array_view::<f64>()?, b.to_plain_array_view::<f64>()?);
1060        let flat = at.unwrap_or(0);
1061        bail!(
1062            "Mismatch. Max ULP distance ({dt:?}): {} > {}, at {} ({} != {}).",
1063            worst,
1064            max_ulp,
1065            indices,
1066            a.iter().nth(flat).copied().unwrap_or(f64::NAN),
1067            b.iter().nth(flat).copied().unwrap_or(f64::NAN),
1068        );
1069    }
1070
1071    /// Transform the tensor into a `ndarray::Array`.
1072    pub fn into_plain_array<D: Datum>(self) -> TractResult<ArrayD<D>> {
1073        Ok(self.to_plain_array_view::<D>()?.to_owned())
1074    }
1075
1076    /// Transform the tensor into a `ndarray::Array`.
1077    pub unsafe fn into_array_unchecked<D: Datum>(self) -> ArrayD<D> {
1078        unsafe { self.to_array_view_unchecked::<D>().to_owned() }
1079    }
1080
1081    /// Returns a plain array view of the tensor.
1082    ///
1083    /// Errors if the storage is not plain or the datum type does not match `D`.
1084    #[inline]
1085    pub fn to_plain_array_view<D: Datum>(&self) -> TractResult<ArrayViewD<'_, D>> {
1086        self.try_as_plain()?.to_array_view::<D>()
1087    }
1088
1089    /// Returns a mutable plain array view of the tensor.
1090    ///
1091    /// Errors if the storage is not plain or the datum type does not match `D`.
1092    #[inline]
1093    pub fn to_plain_array_view_mut<D: Datum>(&mut self) -> TractResult<ArrayViewMutD<'_, D>> {
1094        self.check_for_access::<D>()?;
1095        ensure!(self.storage.as_plain_mut().is_some(), "Tensor storage is not plain");
1096        unsafe { Ok(self.to_array_view_mut_unchecked()) }
1097    }
1098
1099    fn check_for_access<D: Datum>(&self) -> TractResult<()> {
1100        ensure!(
1101            self.datum_type().unquantized() == D::datum_type().unquantized(),
1102            "Tensor datum type error: tensor is {:?}, accessed as {:?}",
1103            self.datum_type(),
1104            D::datum_type(),
1105        );
1106        Ok(())
1107    }
1108
1109    /// Transform the data as a `ndarray::Array`.
1110    pub unsafe fn to_array_view_unchecked<D: Datum>(&self) -> ArrayViewD<'_, D> {
1111        if self.len() != 0 {
1112            unsafe {
1113                ArrayViewD::from_shape_ptr(&*self.shape, self.plain_storage().as_ptr() as *const D)
1114            }
1115        } else {
1116            ArrayViewD::from_shape(&*self.shape, &[]).unwrap()
1117        }
1118    }
1119
1120    /// Transform the data as a mutable `ndarray::Array`.
1121    pub unsafe fn to_array_view_mut_unchecked<D: Datum>(&mut self) -> ArrayViewMutD<'_, D> {
1122        if self.len() != 0 {
1123            unsafe {
1124                let ptr = self.plain_storage_mut().as_mut_ptr() as *mut D;
1125                ArrayViewMutD::from_shape_ptr(&*self.shape, ptr)
1126            }
1127        } else {
1128            ArrayViewMutD::from_shape(&*self.shape, &mut []).unwrap()
1129        }
1130    }
1131
1132    /// Access the data as a pointer.
1133    pub fn as_ptr<D: Datum>(&self) -> TractResult<*const D> {
1134        self.check_for_access::<D>()?;
1135        Ok(self.plain_storage().as_ptr() as *const D)
1136    }
1137
1138    /// Access the data as a pointer.
1139    pub unsafe fn as_ptr_unchecked<D: Datum>(&self) -> *const D {
1140        self.plain_storage().as_ptr() as *const D
1141    }
1142
1143    /// Access the data as a pointer.
1144    pub unsafe fn as_ptr_mut_unchecked<D: Datum>(&mut self) -> *mut D {
1145        self.plain_storage_mut().as_mut_ptr() as *mut D
1146    }
1147
1148    /// Access the data as a mutable pointer.
1149    pub fn as_ptr_mut<D: Datum>(&mut self) -> TractResult<*mut D> {
1150        self.as_ptr::<D>().map(|p| p as *mut D)
1151    }
1152
1153    /// Access the data as a slice.
1154    pub unsafe fn as_slice_unchecked<D: Datum>(&self) -> &[D] {
1155        if self.storage.byte_len() == 0 {
1156            &[]
1157        } else {
1158            unsafe { std::slice::from_raw_parts::<D>(self.as_ptr_unchecked(), self.len()) }
1159        }
1160    }
1161
1162    /// Access the data as a mutable slice.
1163    pub unsafe fn as_slice_mut_unchecked<D: Datum>(&mut self) -> &mut [D] {
1164        if self.storage.byte_len() == 0 {
1165            &mut []
1166        } else {
1167            unsafe { std::slice::from_raw_parts_mut::<D>(self.as_ptr_mut_unchecked(), self.len()) }
1168        }
1169    }
1170
1171    /// Make the tensor a scalar tensor (assumes it contains a single value).
1172    pub fn to_scalar_tensor(&self) -> TractResult<Tensor> {
1173        fn to_scalar_tensor_t<D: Datum>(t: &Tensor) -> TractResult<Tensor> {
1174            Ok(litteral::tensor0(t.try_as_plain()?.to_scalar::<D>()?.clone()))
1175        }
1176        dispatch_datum!(to_scalar_tensor_t(self.datum_type())(self))
1177    }
1178
1179    /// Access the data as a scalar.
1180    pub unsafe fn to_scalar_unchecked<D: Datum>(&self) -> &D {
1181        unsafe { &*(self.plain_storage().as_ptr() as *const D) }
1182    }
1183
1184    /// Mutable access the data as a scalar.
1185    pub fn to_scalar_mut<D: Datum>(&mut self) -> TractResult<&mut D> {
1186        self.check_for_access::<D>()?;
1187        if self.len() == 0 {
1188            bail!("to_scalar_mut called on empty tensor ({:?})", self)
1189        }
1190        if self.len() > 1 {
1191            bail!("to_scalar called on a tensor with multiple values ({:?})", self)
1192        }
1193        unsafe { Ok(self.to_scalar_mut_unchecked()) }
1194    }
1195
1196    /// Mutable access the data as a scalar.
1197    pub unsafe fn to_scalar_mut_unchecked<D: Datum>(&mut self) -> &mut D {
1198        unsafe { &mut *(self.plain_storage_mut().as_mut_ptr() as *mut D) }
1199    }
1200
1201    pub fn as_bytes(&self) -> &[u8] {
1202        self.plain_storage().as_bytes()
1203    }
1204
1205    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
1206        self.plain_storage_mut().as_bytes_mut()
1207    }
1208
1209    unsafe fn is_uniform_t<T: Datum>(&self) -> bool {
1210        let slice = unsafe { self.as_slice_unchecked::<T>() };
1211        slice[1..].iter().all(|x| x == &slice[0])
1212    }
1213
1214    pub fn is_uniform(&self) -> bool {
1215        if self.is_exotic() {
1216            return false;
1217        }
1218        if self.len() <= 1 {
1219            return true;
1220        }
1221        unsafe { dispatch_datum!(Tensor::is_uniform_t(self.datum_type())(self)) }
1222    }
1223
1224    unsafe fn as_uniform_t<T: Datum>(&self) -> Tensor {
1225        let v: T = unsafe { self.as_slice_unchecked::<T>() }[0].clone();
1226        litteral::tensor0(v)
1227    }
1228
1229    pub fn as_uniform(&self) -> Option<Tensor> {
1230        if self.len() >= 1 && self.is_uniform() {
1231            unsafe {
1232                let mut t = dispatch_datum!(Tensor::as_uniform_t(self.datum_type())(self));
1233                t.set_datum_type(self.datum_type());
1234                Some(t)
1235            }
1236        } else {
1237            None
1238        }
1239    }
1240
1241    pub fn is_all_zero(&self) -> TractResult<bool> {
1242        Ok(self.len() == 0 || self.as_uniform().map(|t| t.is_zero().unwrap()).unwrap_or(false))
1243    }
1244
1245    pub fn is_zero(&self) -> TractResult<bool> {
1246        Ok(self == &Tensor::zero_scalar_dt(self.dt)?)
1247    }
1248
1249    unsafe fn natural_cast<
1250        Source: Datum + num_traits::AsPrimitive<Target>,
1251        Target: Datum + Copy,
1252    >(
1253        &self,
1254        other: &mut Tensor,
1255    ) {
1256        unsafe {
1257            self.as_slice_unchecked::<Source>()
1258                .iter()
1259                .zip(other.as_slice_mut_unchecked::<Target>().iter_mut())
1260                .for_each(|(s, d)| *d = s.as_())
1261        };
1262    }
1263
1264    unsafe fn cast_number_to_bool<Source: Datum + num_traits::Zero>(&self, other: &mut Tensor) {
1265        unsafe {
1266            self.as_slice_unchecked::<Source>()
1267                .iter()
1268                .zip(other.as_slice_mut_unchecked::<bool>().iter_mut())
1269                .for_each(|(s, d)| *d = !s.is_zero());
1270        }
1271    }
1272
1273    unsafe fn cast_from_string<Target: Datum + core::str::FromStr>(
1274        &self,
1275        other: &mut Tensor,
1276    ) -> TractResult<()> {
1277        unsafe {
1278            for (s, d) in self
1279                .as_slice_unchecked::<String>()
1280                .iter()
1281                .zip(other.as_slice_mut_unchecked::<Target>().iter_mut())
1282            {
1283                *d = s
1284                    .parse()
1285                    .map_err(|_| format_err!("Can not parse as {:?}", Target::datum_type()))?;
1286            }
1287            Ok(())
1288        }
1289    }
1290
1291    unsafe fn cast_to_string<Source: Datum>(&self, other: &mut Tensor) {
1292        unsafe {
1293            for (s, d) in self
1294                .as_slice_unchecked::<Source>()
1295                .iter()
1296                .zip(other.as_slice_mut_unchecked::<String>().iter_mut())
1297            {
1298                *d = s.to_string()
1299            }
1300        }
1301    }
1302
1303    /// Optionnaly convert data to a tensor for a new DatumType.
1304    pub fn cast_to<D: Datum>(&self) -> TractResult<Cow<'_, Tensor>> {
1305        self.cast_to_dt(D::datum_type())
1306    }
1307
1308    /// Optionnaly convert data to a tensor for a new DatumType.
1309    #[allow(clippy::redundant_closure_call)]
1310    pub fn cast_to_dt(&self, dst_dt: DatumType) -> TractResult<Cow<'_, Tensor>> {
1311        unsafe {
1312            if self.dt == dst_dt {
1313                return Ok(Cow::Borrowed(self));
1314            }
1315            if self.dt == TDim::datum_type() && (dst_dt.is_integer() || dst_dt.is_float()) {
1316                let slice = self.as_slice_unchecked::<TDim>();
1317                let mut ints = Self::uninitialized::<i64>(&self.shape)?;
1318                let ints_slice = ints.as_slice_mut_unchecked::<i64>();
1319                for i in 0..self.len() {
1320                    ints_slice[i] = slice[i].to_i64()?;
1321                }
1322                return Ok(Cow::Owned(ints.cast_to_dt(dst_dt)?.into_owned()));
1323            }
1324            if self.dt == bool::datum_type()
1325                && (dst_dt.is_integer() || dst_dt.is_float() || dst_dt == TDim::datum_type())
1326            {
1327                let slice = self.as_slice_unchecked::<bool>();
1328                let mut ints = Self::uninitialized::<i8>(&self.shape)?;
1329                let ints_slice = ints.as_slice_mut_unchecked::<i8>();
1330                for i in 0..self.len() {
1331                    ints_slice[i] = slice[i] as usize as i8;
1332                }
1333                return Ok(Cow::Owned(ints.cast_to_dt(dst_dt)?.into_owned()));
1334            }
1335            let mut result = Self::uninitialized_dt(dst_dt, &self.shape)?;
1336            if self.dt == DatumType::String {
1337                dispatch_numbers!(Self::cast_from_string(dst_dt)(self, &mut result))?;
1338                return Ok(Cow::Owned(result));
1339            }
1340            if dst_dt == DatumType::String {
1341                dispatch_datum!(Self::cast_to_string(self.dt)(self, &mut result));
1342                return Ok(Cow::Owned(result));
1343            }
1344            macro_rules! n {
1345                ($source:ty) => {
1346                    if <$source>::datum_type() == self.datum_type() {
1347                        match dst_dt {
1348                            DatumType::I8 => self.natural_cast::<$source, i8>(&mut result),
1349                            DatumType::I16 => self.natural_cast::<$source, i16>(&mut result),
1350                            DatumType::I32 => self.natural_cast::<$source, i32>(&mut result),
1351                            DatumType::I64 => self.natural_cast::<$source, i64>(&mut result),
1352                            DatumType::U8 => self.natural_cast::<$source, u8>(&mut result),
1353                            DatumType::U16 => self.natural_cast::<$source, u16>(&mut result),
1354                            DatumType::U32 => self.natural_cast::<$source, u32>(&mut result),
1355                            DatumType::U64 => self.natural_cast::<$source, u64>(&mut result),
1356                            DatumType::F16 => self.natural_cast::<$source, f16>(&mut result),
1357                            DatumType::F32 => self.natural_cast::<$source, f32>(&mut result),
1358                            DatumType::F64 => self.natural_cast::<$source, f64>(&mut result),
1359                            DatumType::TDim => {
1360                                let ints = self.cast_to::<i32>()?;
1361                                let slice = ints.as_slice_unchecked::<i32>();
1362                                let result = result.as_slice_mut_unchecked::<TDim>();
1363                                for i in 0..self.len() {
1364                                    result[i] = slice[i].into();
1365                                }
1366                            }
1367                            DatumType::Bool => self.cast_number_to_bool::<$source>(&mut result),
1368                            _ => todo!(),
1369                        }
1370                        return Ok(Cow::Owned(result));
1371                    };
1372                };
1373            }
1374            //If there is no quantization
1375            if !dst_dt.is_quantized() && !self.datum_type().is_quantized() {
1376                n!(u8);
1377                n!(u16);
1378                n!(u32);
1379                n!(u64);
1380                n!(i8);
1381                n!(i16);
1382                n!(i32);
1383                n!(i64);
1384                n!(f16);
1385                n!(f32);
1386                n!(f64);
1387            } else {
1388                let (s_zp, s_scale) = self.datum_type().zp_scale();
1389                let (d_zp, d_scale) = dst_dt.zp_scale();
1390                if self.datum_type().is_quantized() && dst_dt.is_float() {
1391                    macro_rules! q_to_fp {
1392                        ($source:ty, $dest:ty) => {
1393                            if <$source>::datum_type().unquantized()
1394                                == self.datum_type().unquantized()
1395                                && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1396                            {
1397                                self.as_slice_unchecked::<$source>()
1398                                    .iter()
1399                                    .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1400                                    .for_each(|(&s, d)| {
1401                                        *d = (s as $dest - s_zp as $dest) * s_scale as $dest;
1402                                    });
1403                                return Ok(Cow::Owned(result));
1404                            }
1405                        };
1406                    }
1407                    q_to_fp!(i8, f64);
1408                    q_to_fp!(i8, f32);
1409                    q_to_fp!(u8, f64);
1410                    q_to_fp!(u8, f32);
1411                }
1412                //TODO: optimize scale_by
1413                macro_rules! q8_to_q8 {
1414                    ($typ:ty) => {
1415                        if dst_dt.unquantized() == <$typ>::datum_type() {
1416                            self.as_slice_unchecked::<$typ>()
1417                                .iter()
1418                                .zip(result.as_slice_mut_unchecked::<$typ>().iter_mut())
1419                                .for_each(|(&s, d)| {
1420                                    *d = (d_zp as i32
1421                                        + scale_by(s as i32 - s_zp as i32, s_scale / d_scale))
1422                                    .clamp_cast()
1423                                });
1424                            return Ok(Cow::Owned(result));
1425                        }
1426                    };
1427                }
1428
1429                macro_rules! q_via_f32 {
1430                    ($source:ty, $dest:ty, $round:expr) => {
1431                        if <$source>::datum_type().unquantized() == self.datum_type().unquantized()
1432                            && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1433                        {
1434                            self.as_slice_unchecked::<$source>()
1435                                .iter()
1436                                .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1437                                .for_each(|(&s, d)| {
1438                                    let s_float = (s as f32 - s_zp as f32) * s_scale as f32;
1439                                    let d_float = s_float as f32 / d_scale as f32 + d_zp as f32;
1440                                    *d = $round(d_float);
1441                                });
1442                            return Ok(Cow::Owned(result));
1443                        }
1444                    };
1445                }
1446
1447                macro_rules! q_n {
1448                    (clamp $source:ty, $dest:ty) => {{
1449                        if <$source>::datum_type().unquantized() == self.datum_type().unquantized()
1450                            && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1451                        {
1452                            self.as_slice_unchecked::<$source>()
1453                                .iter()
1454                                .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1455                                .for_each(|(&s, d)| {
1456                                    *d = s.clamp_cast();
1457                                });
1458                            return Ok(Cow::Owned(result));
1459                        }
1460                    }};
1461                    ($source:ty, $dest:ty) => {{
1462                        if <$source>::datum_type().unquantized() == self.datum_type().unquantized()
1463                            && <$dest>::datum_type().unquantized() == dst_dt.unquantized()
1464                        {
1465                            self.as_slice_unchecked::<$source>()
1466                                .iter()
1467                                .zip(result.as_slice_mut_unchecked::<$dest>().iter_mut())
1468                                .for_each(|(&s, d)| {
1469                                    *d = s as $dest;
1470                                });
1471                            return Ok(Cow::Owned(result));
1472                        }
1473                    }};
1474                }
1475
1476                if dst_dt.unquantized() == self.datum_type().unquantized()
1477                    && dst_dt.is_quantized()
1478                    && self.datum_type().is_quantized()
1479                {
1480                    q8_to_q8!(i8);
1481                    q8_to_q8!(u8);
1482                }
1483
1484                q_via_f32!(f32, i8, |f| round_ties_to_even(f).clamp_cast());
1485                q_via_f32!(f32, u8, |f| round_ties_to_even(f).clamp_cast());
1486                q_via_f32!(f32, i32, |f| round_ties_to_even(f).clamp_cast());
1487                q_via_f32!(i8, f32, |f| f);
1488                q_via_f32!(u8, f32, |f| f);
1489                q_via_f32!(i32, f32, |f| f);
1490
1491                if dst_dt.is_quantized() && self.datum_type().is_quantized() {
1492                    q_via_f32!(u8, i8, |f| round_ties_to_even(f).clamp_cast());
1493                    q_via_f32!(i8, u8, |f| round_ties_to_even(f).clamp_cast());
1494                    q_via_f32!(i32, u8, |f| round_ties_to_even(f).clamp_cast());
1495                    q_via_f32!(i32, i8, |f| round_ties_to_even(f).clamp_cast());
1496                    q_via_f32!(u8, i32, |f| round_ties_to_even(f).clamp_cast());
1497                    q_via_f32!(i8, i32, |f| round_ties_to_even(f).clamp_cast());
1498
1499                    // ensure cast to different scale offset work
1500                    q_via_f32!(i8, i8, |f| round_ties_to_even(f).clamp_cast());
1501                    q_via_f32!(u8, u8, |f| round_ties_to_even(f).clamp_cast());
1502                }
1503
1504                q_n!(i8, i32);
1505                q_n!(i8, u32);
1506                q_n!(u8, i32);
1507                q_n!(u8, u32);
1508                q_n!(clamp i32, i8);
1509                q_n!(clamp i32, u8);
1510                q_n!(clamp u32, i8);
1511                q_n!(clamp u32, u8);
1512                q_n!(i8, i8);
1513                q_n!(u8, u8);
1514                q_n!(i32, i32);
1515                q_n!(u32, u32);
1516            }
1517
1518            bail!("Unsupported cast from {:?} to {:?}", self.dt, dst_dt)
1519        }
1520    }
1521
1522    /// Access the data as a scalar, after a cast.
1523    pub fn cast_to_scalar<D: Datum + Copy>(&self) -> TractResult<D> {
1524        let casted = self.cast_to::<D>()?;
1525        casted.try_as_plain()?.to_scalar::<D>().copied()
1526    }
1527
1528    /// Access the nth element of the tensor, returned as a 0-rank Tensor
1529    pub fn nth(&self, nth: usize) -> TractResult<Tensor> {
1530        if nth >= self.len() {
1531            bail!(
1532                "nth called with {}th element on a tensor of len {} ({:?}",
1533                nth,
1534                self.len(),
1535                self
1536            );
1537        }
1538        unsafe fn nth_t<T: Datum>(me: &Tensor, nth: usize, output: &mut Tensor) {
1539            unsafe {
1540                let value = me.as_slice_unchecked::<T>()[nth].clone();
1541                output.as_slice_mut_unchecked::<T>()[0] = value;
1542            }
1543        }
1544        unsafe {
1545            let mut output = Tensor::uninitialized_dt(self.datum_type(), &[])?;
1546            dispatch_datum_by_size!(nth_t(self.datum_type())(self, nth, &mut output));
1547            Ok(output)
1548        }
1549    }
1550
1551    /// Strict equality test on tensors.
1552    fn eq_dt(&self, other: &Tensor) -> TractResult<bool> {
1553        unsafe fn eq_t<D: Datum>(me: &Tensor, other: &Tensor) -> TractResult<bool> {
1554            unsafe {
1555                if D::datum_type().is_float() {
1556                    return dispatch_floatlike!(float_eq_t(D::datum_type())(me, other));
1557                }
1558                Ok(izip!(me.as_slice_unchecked::<D>(), other.as_slice_unchecked::<D>())
1559                    .all(|(a, b)| a == b))
1560            }
1561        }
1562
1563        unsafe fn float_eq_t<D: Datum + Float>(me: &Tensor, other: &Tensor) -> TractResult<bool> {
1564            unsafe {
1565                Ok(izip!(me.as_slice_unchecked::<D>(), other.as_slice_unchecked::<D>())
1566                    .all(|(a, b)| (a.is_nan() && b.is_nan()) || a == b))
1567            }
1568        }
1569
1570        unsafe {
1571            Ok(self.datum_type() == other.datum_type()
1572                && self.shape() == other.shape()
1573                && dispatch_datum!(eq_t(self.dt)(self, other))?)
1574        }
1575    }
1576
1577    fn from_datum<T: Datum>(mut it: ArrayD<T>) -> Tensor {
1578        unsafe {
1579            let mut t = Self::uninitialized::<T>(it.shape()).unwrap();
1580            if let Some(slice) = it.as_slice_mut() {
1581                if t.datum_type().is_copy() {
1582                    std::ptr::copy_nonoverlapping(
1583                        slice.as_ptr() as *const i8,
1584                        t.as_ptr_mut_unchecked(),
1585                        t.plain_storage().layout().size(),
1586                    );
1587                } else {
1588                    t.as_slice_mut_unchecked::<T>()
1589                        .iter_mut()
1590                        .zip(slice.iter_mut())
1591                        .for_each(|(t, s)| *t = std::mem::take(s));
1592                }
1593                return t;
1594            }
1595            if it.strides().iter().all(|&s| s > 0) && it.as_slice_memory_order().is_some() {
1596                let mut len_and_strides: TVec<(usize, usize)> = tvec!();
1597                for (len, stride) in itertools::izip!(it.shape(), it.strides(), t.strides())
1598                    .sorted_by_key(|(_, src, _)| *src)
1599                    .map(|(l, _, dst)| (*l as isize, *dst))
1600                {
1601                    if !len_and_strides.is_empty()
1602                        && len_and_strides.last().unwrap().1 * len_and_strides.last().unwrap().0
1603                            == stride as usize
1604                    {
1605                        len_and_strides.last_mut().unwrap().0 *= len as usize;
1606                    } else {
1607                        len_and_strides.push((len as usize, stride as usize));
1608                    }
1609                }
1610                len_and_strides.reverse();
1611                crate::scatter::scatter_contig_data(
1612                    it.as_ptr(),
1613                    t.as_ptr_mut_unchecked(),
1614                    &len_and_strides,
1615                );
1616                return t;
1617            }
1618            // finally use ndarray into_iter()
1619            t.as_slice_mut_unchecked().iter_mut().zip(it).for_each(|(t, a)| *t = a);
1620            t
1621        }
1622    }
1623
1624    pub fn deep_clone(&self) -> Tensor {
1625        if self.is_exotic() {
1626            return Tensor {
1627                dt: self.dt,
1628                shape: self.shape.clone(),
1629                strides: self.strides.clone(),
1630                len: self.len,
1631                storage: self.storage.deep_clone(),
1632            };
1633        }
1634        unsafe {
1635            let mut tensor = Tensor::uninitialized_dt(self.datum_type(), self.shape()).unwrap();
1636            if self.len() > 0 {
1637                if self.dt.is_copy() {
1638                    self.plain_storage().as_ptr().copy_to_nonoverlapping(
1639                        tensor.as_bytes_mut().as_mut_ptr(),
1640                        self.plain_storage().layout().size(),
1641                    )
1642                } else if self.dt == DatumType::String {
1643                    tensor
1644                        .as_slice_mut_unchecked::<String>()
1645                        .clone_from_slice(self.as_slice_unchecked());
1646                } else if self.dt == DatumType::Blob {
1647                    tensor
1648                        .as_slice_mut_unchecked::<Blob>()
1649                        .clone_from_slice(self.as_slice_unchecked());
1650                } else if self.dt == DatumType::TDim {
1651                    tensor
1652                        .as_slice_mut_unchecked::<TDim>()
1653                        .clone_from_slice(self.as_slice_unchecked());
1654                }
1655            }
1656            tensor
1657        }
1658    }
1659
1660    pub fn slice(&self, axis: usize, start: usize, end: usize) -> TractResult<Tensor> {
1661        if axis >= self.rank() {
1662            bail!("Can not slice at axis {} tensor {:?}", axis, self);
1663        }
1664        if start > self.shape[axis] || end > self.shape[axis] || start >= end {
1665            bail!("Invalid slicing range {start}..{end} on axis {axis} for {self:?}");
1666        }
1667        fn slice_t<T: Datum>(
1668            t: &Tensor,
1669            axis: usize,
1670            start: usize,
1671            end: usize,
1672        ) -> TractResult<Tensor> {
1673            Ok(t.to_plain_array_view::<T>()?
1674                .slice_axis(ndarray::Axis(axis), (start..end).into())
1675                .into_owned()
1676                .into_tensor())
1677        }
1678        dispatch_datum!(slice_t(self.datum_type())(self, axis, start, end))
1679    }
1680
1681    #[inline]
1682    pub fn view(&self) -> view::TensorView<'_> {
1683        unsafe { view::TensorView::view(self) }
1684    }
1685
1686    #[inline]
1687    pub fn view_at_prefix(&self, prefix: &[usize]) -> TractResult<view::TensorView<'_>> {
1688        view::TensorView::at_prefix(self, prefix)
1689    }
1690
1691    #[inline]
1692    pub fn view_offsetting(&self, coords: &[usize]) -> TractResult<view::TensorView<'_>> {
1693        view::TensorView::offsetting(self, coords)
1694    }
1695
1696    #[inline]
1697    pub unsafe fn view_offsetting_unchecked(&self, coords: &[usize]) -> view::TensorView<'_> {
1698        unsafe { view::TensorView::offsetting_unchecked(self, coords) }
1699    }
1700
1701    #[inline]
1702    pub fn view_mut(&mut self) -> view::TensorView<'_> {
1703        unsafe { view::TensorView::view(self) }
1704    }
1705
1706    #[inline]
1707    pub fn view_at_prefix_mut(&mut self, prefix: &[usize]) -> TractResult<view::TensorView<'_>> {
1708        view::TensorView::at_prefix(self, prefix)
1709    }
1710
1711    #[inline]
1712    pub fn view_offsetting_mut(&mut self, coords: &[usize]) -> TractResult<view::TensorView<'_>> {
1713        view::TensorView::offsetting(self, coords)
1714    }
1715
1716    /// Offsets the tensor as an i8 type if it's an u8 type, otherwise passes it unchanged.
1717    pub fn offset_u8_as_i8(self: &Arc<Self>) -> Arc<Self> {
1718        let mut t = if let DatumType::U8 = self.dt.unquantized() {
1719            self.try_as_plain()
1720                .unwrap()
1721                .to_array_view::<u8>()
1722                .unwrap()
1723                .mapv(|v| v.wrapping_sub(128) as i8)
1724                .into_tensor()
1725        } else {
1726            return self.clone();
1727        };
1728
1729        if let DatumType::QU8(qp) = self.dt {
1730            if let QParams::ZpScale { zero_point, scale } = qp {
1731                t.dt = DatumType::QI8(QParams::ZpScale { zero_point: zero_point - 128, scale });
1732            } else {
1733                t.dt = DatumType::QI8(qp);
1734            }
1735        }
1736
1737        t.into_arc_tensor()
1738    }
1739
1740    /// Offsets the tensor as an u8 type if it's an i8 type, otherwise passes it unchanged.
1741    pub fn offset_i8_as_u8(self: &Arc<Self>) -> Arc<Self> {
1742        let mut t = if let DatumType::I8 = self.dt.unquantized() {
1743            self.try_as_plain()
1744                .unwrap()
1745                .to_array_view::<i8>()
1746                .unwrap()
1747                .mapv(|v| (v as u8).wrapping_add(128))
1748                .into_tensor()
1749        } else {
1750            return self.clone();
1751        };
1752
1753        if let DatumType::QI8(qp) = self.dt {
1754            if let QParams::ZpScale { zero_point, scale } = qp {
1755                t.dt = DatumType::QU8(QParams::ZpScale { zero_point: zero_point + 128, scale });
1756            } else {
1757                t.dt = DatumType::QU8(qp);
1758            }
1759        }
1760        t.into_arc_tensor()
1761    }
1762
1763    pub fn to_aligned_default(&self) -> TractResult<Self> {
1764        if self.dt.is_copy() {
1765            unsafe {
1766                let mut t = Self::uninitialized_dt(self.dt, &self.shape)?;
1767                t.as_bytes_mut().copy_from_slice(self.as_bytes());
1768                Ok(t)
1769            }
1770        } else {
1771            let mut t = Self::zero_dt(self.dt, &self.shape)?;
1772            if self.dt == String::datum_type() {
1773                t.try_as_plain_mut()?
1774                    .as_slice_mut::<String>()?
1775                    .clone_from_slice(self.try_as_plain()?.as_slice()?);
1776            } else if self.dt == Blob::datum_type() {
1777                t.try_as_plain_mut()?
1778                    .as_slice_mut::<Blob>()?
1779                    .clone_from_slice(self.try_as_plain()?.as_slice()?);
1780            } else if self.dt == TDim::datum_type() {
1781                t.try_as_plain_mut()?
1782                    .as_slice_mut::<TDim>()?
1783                    .clone_from_slice(self.try_as_plain()?.as_slice()?);
1784            }
1785            Ok(t)
1786        }
1787    }
1788
1789    pub fn natural_strides(shape: &[usize]) -> TVec<isize> {
1790        let mut strides = tvec!();
1791        compute_natural_stride_to(&mut strides, shape);
1792        strides
1793    }
1794
1795    pub fn into_blob(mut self) -> TractResult<Blob> {
1796        ensure!(self.dt.is_copy());
1797        let storage =
1798            std::mem::replace(&mut self.storage, StorageKind::Plain(PlainStorage::default()));
1799        Ok(storage.into_plain().context("Storage is not plain")?.into_blob())
1800    }
1801}
1802
1803impl PartialEq for Tensor {
1804    fn eq(&self, other: &Tensor) -> bool {
1805        if self.dt != other.dt || self.shape != other.shape {
1806            return false;
1807        }
1808        match (self.storage.as_plain(), other.storage.as_plain()) {
1809            (Some(_), Some(_)) => self.eq_dt(other).unwrap_or(false),
1810            (None, None) => self.storage == other.storage,
1811            _ => false,
1812        }
1813    }
1814}
1815
1816impl Eq for Tensor {}
1817
1818impl fmt::Debug for Tensor {
1819    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
1820        let content = self.dump(false).unwrap_or_else(|e| format!("Error : {e:?}"));
1821        write!(formatter, "{content}")
1822    }
1823}
1824
1825#[cfg(feature = "complex")]
1826pub fn reinterpret_inner_dim_as_complex(mut t: Tensor) -> TractResult<Tensor> {
1827    ensure!(
1828        t.shape().last() == Some(&2),
1829        "The last dimension in the tensor shape {:?} must be 2",
1830        t.shape()
1831    );
1832    unsafe {
1833        t.shape.pop();
1834        t.set_datum_type(t.datum_type().complexify()?);
1835        t.update_strides_and_len();
1836        Ok(t)
1837    }
1838}
1839
1840#[cfg(feature = "complex")]
1841pub fn reinterpret_complex_as_inner_dim(mut t: Tensor) -> TractResult<Tensor> {
1842    unsafe {
1843        t.shape.push(2);
1844        t.set_datum_type(t.datum_type().decomplexify()?);
1845        t.update_strides_and_len();
1846        Ok(t)
1847    }
1848}
1849
1850pub fn clip_range_bounds(len: usize, range: impl std::ops::RangeBounds<usize>) -> Range<usize> {
1851    use std::ops::Bound;
1852    let start = match range.start_bound() {
1853        Bound::Included(ix) => *ix,
1854        Bound::Excluded(ix) => ix + 1,
1855        Bound::Unbounded => 0,
1856    };
1857    let end = match range.end_bound() {
1858        Bound::Included(ix) => *ix + 1,
1859        Bound::Excluded(ix) => *ix,
1860        Bound::Unbounded => len,
1861    };
1862    start..end
1863}
1864
1865pub fn natural_strides(shape: &[usize]) -> TVec<isize> {
1866    let mut strides = tvec!();
1867    compute_natural_stride_to(&mut strides, shape);
1868    strides
1869}
1870
1871fn compute_natural_stride_to(strides: &mut TVec<isize>, shape: &[usize]) {
1872    match shape.len() {
1873        0 => (),
1874        1 => strides.push(1),
1875        2 => strides.extend_from_slice(&[shape[1] as isize, 1]),
1876        3 => strides.extend_from_slice(&[(shape[1] * shape[2]) as isize, shape[2] as _, 1]),
1877        4 => strides.extend_from_slice(&[
1878            (shape[1] * shape[2] * shape[3]) as isize,
1879            (shape[2] * shape[3]) as _,
1880            shape[3] as _,
1881            1,
1882        ]),
1883        _ => {
1884            strides.push(1);
1885            for dim in shape.as_ref().iter().skip(1).rev() {
1886                let previous = *strides.last().unwrap();
1887                strides.push(previous * *dim as isize)
1888            }
1889            strides.reverse();
1890        }
1891    }
1892}
1893
1894impl<D: ::ndarray::Dimension, T: Datum> From<Array<T, D>> for Tensor {
1895    fn from(it: Array<T, D>) -> Tensor {
1896        Tensor::from_datum(it.into_dyn())
1897    }
1898}
1899
1900/// Convenient conversion to Tensor.
1901pub trait IntoTensor: Sized {
1902    /// Convert Self to a Tensor.
1903    ///
1904    /// May perform a copy
1905    fn into_tensor(self) -> Tensor;
1906}
1907
1908/// Convenient conversion to Arc<Tensor>.
1909pub trait IntoArcTensor: Sized {
1910    /// Convert Self to a Arc<Tensor>.
1911    ///
1912    /// May perform a copy
1913    fn into_arc_tensor(self) -> Arc<Tensor>;
1914}
1915
1916impl<D: ::ndarray::Dimension, T: Datum> IntoTensor for Array<T, D> {
1917    fn into_tensor(self) -> Tensor {
1918        Tensor::from(self)
1919    }
1920}
1921
1922impl<D: ::ndarray::Dimension, T: Datum> IntoArcTensor for Array<T, D> {
1923    fn into_arc_tensor(self) -> Arc<Tensor> {
1924        Arc::new(Tensor::from(self))
1925    }
1926}
1927
1928impl IntoTensor for Tensor {
1929    fn into_tensor(self) -> Tensor {
1930        self
1931    }
1932}
1933
1934impl IntoTensor for Arc<Tensor> {
1935    fn into_tensor(self) -> Tensor {
1936        Arc::try_unwrap(self).unwrap_or_else(|t| (*t).clone())
1937    }
1938}
1939
1940impl IntoArcTensor for Tensor {
1941    fn into_arc_tensor(self) -> Arc<Tensor> {
1942        Arc::new(self)
1943    }
1944}
1945
1946impl IntoArcTensor for Arc<Tensor> {
1947    fn into_arc_tensor(self) -> Arc<Tensor> {
1948        self
1949    }
1950}
1951
1952#[cfg(test)]
1953mod tests {
1954    use crate::dim::SymbolScope;
1955    use crate::prelude::tensor1;
1956
1957    use super::*;
1958    use litteral::tensor0;
1959    use proptest::collection::vec;
1960    use proptest::prelude::*;
1961
1962    // Regression for sonos/tract#2390: from_raw must reject a content length that
1963    // does not match the declared shape rather than panicking in copy_from_slice.
1964    #[test]
1965    fn from_raw_rejects_length_mismatch() {
1966        // shape [2, 3] of f32 needs 24 bytes; supply 12.
1967        let err = unsafe { Tensor::from_raw_dt(f32::datum_type(), &[2, 3], &[0u8; 12]) }
1968            .expect_err("from_raw must reject a short content buffer, not panic");
1969        assert!(err.to_string().contains("does not match shape"), "unexpected error: {err}");
1970        // Too-long content is rejected as well.
1971        assert!(unsafe { Tensor::from_raw_dt(f32::datum_type(), &[2, 3], &[0u8; 32]) }.is_err());
1972        // Exact match still succeeds.
1973        assert!(unsafe { Tensor::from_raw_dt(f32::datum_type(), &[2, 3], &[0u8; 24]) }.is_ok());
1974    }
1975
1976    #[derive(Debug)]
1977    struct PermuteAxisProblem {
1978        shape: Vec<usize>,
1979        permutation: Vec<usize>,
1980    }
1981
1982    impl Arbitrary for PermuteAxisProblem {
1983        type Strategy = BoxedStrategy<PermuteAxisProblem>;
1984        type Parameters = ();
1985
1986        fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
1987            (0..8usize)
1988                .prop_flat_map(|rank| {
1989                    let permute: Vec<usize> = (0..rank).collect();
1990                    (proptest::collection::vec(1..5usize, rank), Just(permute).prop_shuffle())
1991                })
1992                .prop_map(|(shape, permutation)| PermuteAxisProblem { shape, permutation })
1993                .boxed()
1994        }
1995    }
1996
1997    impl PermuteAxisProblem {
1998        fn input(&self) -> ArrayD<i32> {
1999            let mut i = 0;
2000            ArrayD::from_shape_simple_fn(&*self.shape, || {
2001                i += 1;
2002                i
2003            })
2004            .permuted_axes(&*self.permutation)
2005        }
2006
2007        fn reference(&self) -> Tensor {
2008            let values: Vec<i32> = self.input().iter().copied().collect();
2009            let shape = self.permutation.iter().map(|ix| self.shape[*ix]).collect::<TVec<usize>>();
2010            super::litteral::tensor1(&values).into_shape(&shape).unwrap()
2011        }
2012
2013        fn tract(&self) -> Tensor {
2014            Tensor::from(self.input())
2015        }
2016
2017        fn check(&self) -> proptest::test_runner::TestCaseResult {
2018            prop_assert_eq!(self.tract(), self.reference());
2019            Ok(())
2020        }
2021    }
2022
2023    proptest::proptest! {
2024        #[test]
2025        fn prop(pb: PermuteAxisProblem) {
2026            pb.check().unwrap();
2027        }
2028    }
2029
2030    #[test]
2031    fn t_1_2() {
2032        PermuteAxisProblem { shape: vec![2, 1], permutation: vec![1, 0] }.check().unwrap();
2033    }
2034
2035    #[test]
2036    fn t_2_2() {
2037        PermuteAxisProblem { shape: vec![2, 2], permutation: vec![1, 0] }.check().unwrap();
2038    }
2039
2040    #[derive(Debug)]
2041    struct BroadcastVecToShape {
2042        vec: Vec<f32>,
2043        axis: usize,
2044        shape: TVec<usize>,
2045    }
2046
2047    impl BroadcastVecToShape {
2048        fn check(&self) -> proptest::test_runner::TestCaseResult {
2049            let input = tensor1(&self.vec);
2050            let mut intermediate = tvec![1usize; self.shape.len()];
2051            intermediate[self.axis] = self.vec.len();
2052            let reference = input
2053                .clone()
2054                .into_shape(&intermediate)
2055                .unwrap()
2056                .broadcast_to_shape(&self.shape)
2057                .unwrap();
2058            prop_assert_eq!(
2059                reference,
2060                input.broadcast_vector_to_shape(&self.shape, self.axis).unwrap()
2061            );
2062            Ok(())
2063        }
2064    }
2065
2066    impl Arbitrary for BroadcastVecToShape {
2067        type Strategy = BoxedStrategy<BroadcastVecToShape>;
2068        type Parameters = ();
2069
2070        fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
2071            vec(0usize..5, 0usize..4)
2072                .prop_flat_map(|shape| {
2073                    (vec(-10f32..10f32, 0usize..5), Just(shape.clone()), 0..shape.len() + 1)
2074                })
2075                .prop_map(|(vec, mut shape, axis)| {
2076                    shape.insert(axis, vec.len());
2077                    BroadcastVecToShape { vec, shape: shape.into(), axis }
2078                })
2079                .boxed()
2080        }
2081    }
2082
2083    proptest::proptest! {
2084        #[test]
2085        fn broadcast_vector_to_shape_prop(pb: BroadcastVecToShape) {
2086            pb.check().unwrap()
2087        }
2088    }
2089
2090    #[test]
2091    #[cfg(feature = "complex")]
2092    fn test_reinterpret_inner_dim_as_complex() -> TractResult<()> {
2093        let input = crate::internal::tensor2(&[[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]]);
2094        let cplx_input = reinterpret_inner_dim_as_complex(input)?;
2095        let expected = crate::internal::tensor1(&[
2096            Complex::new(1.0f32, 2.0),
2097            Complex::new(3.0, 4.0),
2098            Complex::new(5.0, 6.0),
2099        ]);
2100        assert_eq!(expected, cplx_input);
2101        Ok(())
2102    }
2103
2104    #[test]
2105    #[cfg(feature = "complex")]
2106    fn test_reinterpret_inner_dim_as_complex_2() -> TractResult<()> {
2107        let input =
2108            crate::internal::tensor3(&[[[1i32, 2], [1, 2]], [[3, 4], [3, 4]], [[5, 6], [5, 6]]]);
2109        let cplx_input = reinterpret_inner_dim_as_complex(input)?;
2110        let expected = crate::internal::tensor2(&[
2111            [Complex::new(1i32, 2), Complex::new(1, 2)],
2112            [Complex::new(3, 4), Complex::new(3, 4)],
2113            [Complex::new(5, 6), Complex::new(5, 6)],
2114        ]);
2115        assert_eq!(expected, cplx_input);
2116        Ok(())
2117    }
2118
2119    #[test]
2120    fn clone_tdim_tensor() {
2121        let symbols = SymbolScope::default();
2122        let a = symbols.sym("a");
2123        let t = tensor0(TDim::from(a));
2124        let _ = t.clone();
2125    }
2126
2127    #[test]
2128    fn ulp_approximation_accepts_within_bound() -> TractResult<()> {
2129        let a = tensor1(&[1.0f32, 2.0, 3.0]);
2130        let b = tensor1(&[
2131            f32::from_bits(1.0f32.to_bits() + 1),
2132            2.0,
2133            f32::from_bits(3.0f32.to_bits() + 2),
2134        ]);
2135        a.close_enough(&b, Approximation::Ulp(2))?;
2136        assert!(a.close_enough(&b, Approximation::Ulp(1)).is_err());
2137        assert_eq!(a.max_ulp_distance(&b)?, (2, Some(2)));
2138        Ok(())
2139    }
2140
2141    #[test]
2142    fn ulp_approximation_uses_the_tensor_own_float_type() -> TractResult<()> {
2143        // One f16 rounding step is ~8192 f32 steps. Measuring in f32 would make an
2144        // adjacent-f16 pair look wildly off, so the comparison must stay in f16.
2145        let one = f16::from_f32(1.0);
2146        let a = tensor1(&[one]);
2147        let b = tensor1(&[f16::from_bits(one.to_bits() + 1)]);
2148        assert_eq!(a.ulp_comparison_dt(), DatumType::F16);
2149        assert_eq!(a.max_ulp_distance(&b)?, (1, Some(0)));
2150        a.close_enough(&b, Approximation::Ulp(1))?;
2151        Ok(())
2152    }
2153
2154    #[test]
2155    fn ulp_approximation_is_scale_free() -> TractResult<()> {
2156        // The same relative error at wildly different magnitudes reads the same,
2157        // which a shared atol cannot do.
2158        let a = tensor1(&[1e-30f32, 1e30]);
2159        let b = tensor1(&[
2160            f32::from_bits(1e-30f32.to_bits() + 1),
2161            f32::from_bits(1e30f32.to_bits() + 1),
2162        ]);
2163        a.close_enough(&b, Approximation::Ulp(1))?;
2164        Ok(())
2165    }
2166
2167    #[test]
2168    fn ulp_approximation_rejects_shape_mismatch() {
2169        let a = tensor1(&[1.0f32, 2.0]);
2170        let b = tensor1(&[1.0f32]);
2171        assert!(a.close_enough(&b, Approximation::Ulp(1000)).is_err());
2172    }
2173
2174    #[test]
2175    fn ulp_bounds_are_distinguished_by_equality() {
2176        assert_eq!(Approximation::Ulp(1), Approximation::Ulp(1));
2177        assert_ne!(Approximation::Ulp(1), Approximation::Ulp(2));
2178    }
2179}