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