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