1use crate::{DType, PixelFormat, Tensor, TensorMemory, TensorTrait};
5use half::f16;
6use std::fmt;
7
8#[non_exhaustive]
10pub enum TensorDyn {
11 U8(Tensor<u8>),
13 I8(Tensor<i8>),
15 U16(Tensor<u16>),
17 I16(Tensor<i16>),
19 U32(Tensor<u32>),
21 I32(Tensor<i32>),
23 U64(Tensor<u64>),
25 I64(Tensor<i64>),
27 F16(Tensor<f16>),
29 F32(Tensor<f32>),
31 F64(Tensor<f64>),
33}
34
35macro_rules! dispatch {
37 ($self:expr, $method:ident $(, $arg:expr)*) => {
38 match $self {
39 TensorDyn::U8(t) => t.$method($($arg),*),
40 TensorDyn::I8(t) => t.$method($($arg),*),
41 TensorDyn::U16(t) => t.$method($($arg),*),
42 TensorDyn::I16(t) => t.$method($($arg),*),
43 TensorDyn::U32(t) => t.$method($($arg),*),
44 TensorDyn::I32(t) => t.$method($($arg),*),
45 TensorDyn::U64(t) => t.$method($($arg),*),
46 TensorDyn::I64(t) => t.$method($($arg),*),
47 TensorDyn::F16(t) => t.$method($($arg),*),
48 TensorDyn::F32(t) => t.$method($($arg),*),
49 TensorDyn::F64(t) => t.$method($($arg),*),
50 }
51 };
52}
53
54macro_rules! dyn_fanout {
58 ($self:expr, $method:ident $(, $arg:expr)*) => {
59 match $self {
60 TensorDyn::U8(t) => t.$method($($arg),*).map(TensorDyn::U8),
61 TensorDyn::I8(t) => t.$method($($arg),*).map(TensorDyn::I8),
62 TensorDyn::U16(t) => t.$method($($arg),*).map(TensorDyn::U16),
63 TensorDyn::I16(t) => t.$method($($arg),*).map(TensorDyn::I16),
64 TensorDyn::U32(t) => t.$method($($arg),*).map(TensorDyn::U32),
65 TensorDyn::I32(t) => t.$method($($arg),*).map(TensorDyn::I32),
66 TensorDyn::U64(t) => t.$method($($arg),*).map(TensorDyn::U64),
67 TensorDyn::I64(t) => t.$method($($arg),*).map(TensorDyn::I64),
68 TensorDyn::F16(t) => t.$method($($arg),*).map(TensorDyn::F16),
69 TensorDyn::F32(t) => t.$method($($arg),*).map(TensorDyn::F32),
70 TensorDyn::F64(t) => t.$method($($arg),*).map(TensorDyn::F64),
71 }
72 };
73}
74
75macro_rules! downcast_methods {
77 ($variant:ident, $ty:ty, $as_name:ident, $as_mut_name:ident, $into_name:ident) => {
78 pub fn $as_name(&self) -> Option<&Tensor<$ty>> {
80 match self {
81 Self::$variant(t) => Some(t),
82 _ => None,
83 }
84 }
85
86 pub fn $as_mut_name(&mut self) -> Option<&mut Tensor<$ty>> {
88 match self {
89 Self::$variant(t) => Some(t),
90 _ => None,
91 }
92 }
93
94 #[allow(clippy::result_large_err)]
97 pub fn $into_name(self) -> Result<Tensor<$ty>, Self> {
98 match self {
99 Self::$variant(t) => Ok(t),
100 other => Err(other),
101 }
102 }
103 };
104}
105
106impl TensorDyn {
107 pub fn dtype(&self) -> DType {
109 match self {
110 Self::U8(_) => DType::U8,
111 Self::I8(_) => DType::I8,
112 Self::U16(_) => DType::U16,
113 Self::I16(_) => DType::I16,
114 Self::U32(_) => DType::U32,
115 Self::I32(_) => DType::I32,
116 Self::U64(_) => DType::U64,
117 Self::I64(_) => DType::I64,
118 Self::F16(_) => DType::F16,
119 Self::F32(_) => DType::F32,
120 Self::F64(_) => DType::F64,
121 }
122 }
123
124 pub fn shape(&self) -> &[usize] {
126 dispatch!(self, shape)
127 }
128
129 pub fn name(&self) -> String {
131 dispatch!(self, name)
132 }
133
134 pub fn format(&self) -> Option<PixelFormat> {
136 dispatch!(self, format)
137 }
138
139 pub fn width(&self) -> Option<usize> {
141 dispatch!(self, width)
142 }
143
144 pub fn height(&self) -> Option<usize> {
146 dispatch!(self, height)
147 }
148
149 pub fn size(&self) -> usize {
151 dispatch!(self, size)
152 }
153
154 pub fn memory(&self) -> TensorMemory {
156 dispatch!(self, memory)
157 }
158
159 pub fn reshape(&mut self, shape: &[usize]) -> crate::Result<()> {
161 dispatch!(self, reshape, shape)
162 }
163
164 pub fn set_format(&mut self, format: PixelFormat) -> crate::Result<()> {
182 dispatch!(self, set_format, format)
183 }
184
185 pub fn with_format(mut self, format: PixelFormat) -> crate::Result<Self> {
202 self.set_format(format)?;
203 Ok(self)
204 }
205
206 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
208 dispatch!(self, colorimetry)
209 }
210
211 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
213 dispatch!(self, set_colorimetry, c)
214 }
215
216 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
218 self.set_colorimetry(Some(c));
219 self
220 }
221
222 pub fn row_stride(&self) -> Option<usize> {
224 dispatch!(self, row_stride)
225 }
226
227 pub fn effective_row_stride(&self) -> Option<usize> {
229 dispatch!(self, effective_row_stride)
230 }
231
232 pub fn configure_image(
235 &mut self,
236 width: usize,
237 height: usize,
238 format: PixelFormat,
239 ) -> crate::Result<()> {
240 dispatch!(self, configure_image, width, height, format)
241 }
242
243 pub fn set_row_stride(&mut self, stride: usize) -> crate::Result<()> {
249 dispatch!(self, set_row_stride, stride)
250 }
251
252 pub fn with_row_stride(mut self, stride: usize) -> crate::Result<Self> {
254 self.set_row_stride(stride)?;
255 Ok(self)
256 }
257
258 pub fn plane_offset(&self) -> Option<usize> {
260 dispatch!(self, plane_offset)
261 }
262
263 pub fn view_origin(&self) -> Option<crate::ViewOrigin> {
267 dispatch!(self, view_origin)
268 }
269
270 pub fn set_plane_offset(&mut self, offset: usize) {
272 dispatch!(self, set_plane_offset, offset)
273 }
274
275 pub fn batch(&self, n: usize) -> crate::Result<TensorDyn> {
278 dyn_fanout!(self, batch, n)
279 }
280
281 pub fn view(&self, region: crate::Region) -> crate::Result<TensorDyn> {
284 dyn_fanout!(self, view, region)
285 }
286
287 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
292 dispatch!(self, cuda)
293 }
294
295 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
318 dispatch!(self, cuda_map)
319 }
320
321 pub fn quantization(&self) -> Option<&crate::Quantization> {
325 match self {
326 Self::U8(t) => t.quantization(),
327 Self::I8(t) => t.quantization(),
328 Self::U16(t) => t.quantization(),
329 Self::I16(t) => t.quantization(),
330 Self::U32(t) => t.quantization(),
331 Self::I32(t) => t.quantization(),
332 Self::U64(t) => t.quantization(),
333 Self::I64(t) => t.quantization(),
334 Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
335 }
336 }
337
338 pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
342 match self {
343 Self::U8(t) => t.set_quantization(q),
344 Self::I8(t) => t.set_quantization(q),
345 Self::U16(t) => t.set_quantization(q),
346 Self::I16(t) => t.set_quantization(q),
347 Self::U32(t) => t.set_quantization(q),
348 Self::I32(t) => t.set_quantization(q),
349 Self::U64(t) => t.set_quantization(q),
350 Self::I64(t) => t.set_quantization(q),
351 Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
352 field: "dtype_is_integer",
353 expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
354 got: format!("{:?}", self.dtype()),
355 }),
356 }
357 }
358
359 pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
362 self.set_quantization(q)?;
363 Ok(self)
364 }
365
366 pub fn clear_quantization(&mut self) {
368 match self {
369 Self::U8(t) => t.clear_quantization(),
370 Self::I8(t) => t.clear_quantization(),
371 Self::U16(t) => t.clear_quantization(),
372 Self::I16(t) => t.clear_quantization(),
373 Self::U32(t) => t.clear_quantization(),
374 Self::I32(t) => t.clear_quantization(),
375 Self::U64(t) => t.clear_quantization(),
376 Self::I64(t) => t.clear_quantization(),
377 Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
378 }
379 }
380
381 #[cfg(unix)]
383 pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
384 dispatch!(self, clone_fd)
385 }
386
387 #[cfg(target_os = "linux")]
398 pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
399 if self.memory() != TensorMemory::Dma {
400 return Err(crate::Error::NotImplemented(format!(
401 "dmabuf_clone requires DMA-backed tensor, got {:?}",
402 self.memory()
403 )));
404 }
405 self.clone_fd()
406 }
407
408 #[cfg(target_os = "linux")]
419 pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
420 dispatch!(self, dmabuf)
421 }
422
423 pub fn is_multiplane(&self) -> bool {
425 dispatch!(self, is_multiplane)
426 }
427
428 pub fn buffer_identity(&self) -> &crate::BufferIdentity {
438 dispatch!(self, buffer_identity)
439 }
440
441 pub fn aliases(&self, other: &Self) -> bool {
459 if self.buffer_identity().id() == other.buffer_identity().id() {
460 return true;
461 }
462 if self.memory() != other.memory() {
463 return false;
464 }
465 #[cfg(target_os = "linux")]
466 if self.memory() == TensorMemory::Dma {
467 use std::os::fd::AsRawFd;
468 if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
469 return a.as_raw_fd() == b.as_raw_fd();
470 }
471 }
472 false
473 }
474
475 downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
478 downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
479 downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
480 downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
481 downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
482 downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
483 downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
484 downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
485 downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
486 downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
487 downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
488
489 pub fn new(
491 shape: &[usize],
492 dtype: DType,
493 memory: Option<TensorMemory>,
494 name: Option<&str>,
495 ) -> crate::Result<Self> {
496 match dtype {
497 DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
498 DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
499 DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
500 DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
501 DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
502 DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
503 DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
504 DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
505 DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
506 DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
507 DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
508 }
509 }
510
511 #[cfg(unix)]
513 pub fn from_fd(
514 fd: std::os::fd::OwnedFd,
515 shape: &[usize],
516 dtype: DType,
517 name: Option<&str>,
518 ) -> crate::Result<Self> {
519 match dtype {
520 DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
521 DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
522 DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
523 DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
524 DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
525 DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
526 DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
527 DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
528 DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
529 DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
530 DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
531 }
532 }
533
534 pub unsafe fn from_foreign_ptr(
547 ptr: *mut u8,
548 shape: &[usize],
549 dtype: DType,
550 owner: Option<crate::ForeignOwner>,
551 name: Option<&str>,
552 ) -> crate::Result<Self> {
553 match dtype {
554 DType::U8 => Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8),
555 DType::I8 => Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8),
556 DType::U16 => {
557 Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
558 }
559 DType::I16 => {
560 Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
561 }
562 DType::U32 => {
563 Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
564 }
565 DType::I32 => {
566 Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
567 }
568 DType::U64 => {
569 Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
570 }
571 DType::I64 => {
572 Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
573 }
574 DType::F16 => {
575 Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
576 }
577 DType::F32 => {
578 Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
579 }
580 DType::F64 => {
581 Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
582 }
583 }
584 }
585
586 #[cfg(target_os = "macos")]
594 pub unsafe fn from_iosurface(
595 surface_ref: *mut std::ffi::c_void,
596 shape: &[usize],
597 dtype: DType,
598 name: Option<&str>,
599 ) -> crate::Result<Self> {
600 unsafe {
601 match dtype {
602 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
603 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
604 DType::U16 => {
605 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
606 }
607 DType::I16 => {
608 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
609 }
610 DType::U32 => {
611 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
612 }
613 DType::I32 => {
614 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
615 }
616 DType::U64 => {
617 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
618 }
619 DType::I64 => {
620 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
621 }
622 DType::F16 => {
623 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
624 }
625 DType::F32 => {
626 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
627 }
628 DType::F64 => {
629 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
630 }
631 }
632 }
633 }
634
635 #[cfg(target_os = "macos")]
638 pub fn iosurface_id(&self) -> Option<u32> {
639 dispatch!(self, iosurface_id)
640 }
641
642 #[cfg(target_os = "macos")]
646 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
647 dispatch!(self, iosurface_ref)
648 }
649
650 #[cfg(target_os = "macos")]
655 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
656 dispatch!(self, iosurface_physical_dims)
657 }
658
659 pub fn image(
677 width: usize,
678 height: usize,
679 format: PixelFormat,
680 dtype: DType,
681 memory: Option<TensorMemory>,
682 ) -> crate::Result<Self> {
683 match dtype {
684 DType::U8 => Tensor::<u8>::image(width, height, format, memory).map(Self::U8),
685 DType::I8 => Tensor::<i8>::image(width, height, format, memory).map(Self::I8),
686 DType::U16 => Tensor::<u16>::image(width, height, format, memory).map(Self::U16),
687 DType::I16 => Tensor::<i16>::image(width, height, format, memory).map(Self::I16),
688 DType::U32 => Tensor::<u32>::image(width, height, format, memory).map(Self::U32),
689 DType::I32 => Tensor::<i32>::image(width, height, format, memory).map(Self::I32),
690 DType::U64 => Tensor::<u64>::image(width, height, format, memory).map(Self::U64),
691 DType::I64 => Tensor::<i64>::image(width, height, format, memory).map(Self::I64),
692 DType::F16 => Tensor::<f16>::image(width, height, format, memory).map(Self::F16),
693 DType::F32 => Tensor::<f32>::image(width, height, format, memory).map(Self::F32),
694 DType::F64 => Tensor::<f64>::image(width, height, format, memory).map(Self::F64),
695 }
696 }
697
698 pub fn image_with_stride(
724 width: usize,
725 height: usize,
726 format: PixelFormat,
727 dtype: DType,
728 row_stride_bytes: usize,
729 memory: Option<TensorMemory>,
730 ) -> crate::Result<Self> {
731 match dtype {
732 DType::U8 => {
733 Tensor::<u8>::image_with_stride(width, height, format, row_stride_bytes, memory)
734 .map(Self::U8)
735 }
736 DType::I8 => {
737 Tensor::<i8>::image_with_stride(width, height, format, row_stride_bytes, memory)
738 .map(Self::I8)
739 }
740 DType::U16 => {
741 Tensor::<u16>::image_with_stride(width, height, format, row_stride_bytes, memory)
742 .map(Self::U16)
743 }
744 DType::I16 => {
745 Tensor::<i16>::image_with_stride(width, height, format, row_stride_bytes, memory)
746 .map(Self::I16)
747 }
748 DType::U32 => {
749 Tensor::<u32>::image_with_stride(width, height, format, row_stride_bytes, memory)
750 .map(Self::U32)
751 }
752 DType::I32 => {
753 Tensor::<i32>::image_with_stride(width, height, format, row_stride_bytes, memory)
754 .map(Self::I32)
755 }
756 DType::U64 => {
757 Tensor::<u64>::image_with_stride(width, height, format, row_stride_bytes, memory)
758 .map(Self::U64)
759 }
760 DType::I64 => {
761 Tensor::<i64>::image_with_stride(width, height, format, row_stride_bytes, memory)
762 .map(Self::I64)
763 }
764 DType::F16 => {
765 Tensor::<f16>::image_with_stride(width, height, format, row_stride_bytes, memory)
766 .map(Self::F16)
767 }
768 DType::F32 => {
769 Tensor::<f32>::image_with_stride(width, height, format, row_stride_bytes, memory)
770 .map(Self::F32)
771 }
772 DType::F64 => {
773 Tensor::<f64>::image_with_stride(width, height, format, row_stride_bytes, memory)
774 .map(Self::F64)
775 }
776 }
777 }
778}
779
780impl From<Tensor<u8>> for TensorDyn {
783 fn from(t: Tensor<u8>) -> Self {
784 Self::U8(t)
785 }
786}
787
788impl From<Tensor<i8>> for TensorDyn {
789 fn from(t: Tensor<i8>) -> Self {
790 Self::I8(t)
791 }
792}
793
794impl From<Tensor<u16>> for TensorDyn {
795 fn from(t: Tensor<u16>) -> Self {
796 Self::U16(t)
797 }
798}
799
800impl From<Tensor<i16>> for TensorDyn {
801 fn from(t: Tensor<i16>) -> Self {
802 Self::I16(t)
803 }
804}
805
806impl From<Tensor<u32>> for TensorDyn {
807 fn from(t: Tensor<u32>) -> Self {
808 Self::U32(t)
809 }
810}
811
812impl From<Tensor<i32>> for TensorDyn {
813 fn from(t: Tensor<i32>) -> Self {
814 Self::I32(t)
815 }
816}
817
818impl From<Tensor<u64>> for TensorDyn {
819 fn from(t: Tensor<u64>) -> Self {
820 Self::U64(t)
821 }
822}
823
824impl From<Tensor<i64>> for TensorDyn {
825 fn from(t: Tensor<i64>) -> Self {
826 Self::I64(t)
827 }
828}
829
830impl From<Tensor<f16>> for TensorDyn {
831 fn from(t: Tensor<f16>) -> Self {
832 Self::F16(t)
833 }
834}
835
836impl From<Tensor<f32>> for TensorDyn {
837 fn from(t: Tensor<f32>) -> Self {
838 Self::F32(t)
839 }
840}
841
842impl From<Tensor<f64>> for TensorDyn {
843 fn from(t: Tensor<f64>) -> Self {
844 Self::F64(t)
845 }
846}
847
848impl fmt::Debug for TensorDyn {
849 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
850 dispatch!(self, fmt, f)
851 }
852}
853
854#[cfg(test)]
855mod tests {
856 use super::*;
857
858 #[test]
859 fn from_typed_tensor() {
860 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
861 let dyn_t: TensorDyn = t.into();
862 assert_eq!(dyn_t.dtype(), DType::U8);
863 assert_eq!(dyn_t.shape(), &[10]);
864 }
865
866 #[test]
867 fn from_foreign_ptr_wraps_borrowed_memory() {
868 use crate::TensorMapTrait;
869 let mut vec: Vec<f32> = vec![0.0; 4];
872 let ptr = vec.as_mut_ptr() as *mut u8;
873 let owner: crate::ForeignOwner = Box::new(vec);
874 let t = unsafe {
875 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
876 }
877 .unwrap();
878 assert_eq!(t.dtype(), DType::F32);
879 assert_eq!(t.memory(), TensorMemory::Mem);
880 assert_eq!(t.shape(), &[2, 2]);
881 {
882 let mut m = t.as_f32().unwrap().map().unwrap();
883 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
884 }
885 let m = t.as_f32().unwrap().map().unwrap();
886 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
887 }
888
889 #[test]
898 fn from_foreign_ptr_rejects_null_ptr() {
899 let err = unsafe {
900 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
901 }
902 .unwrap_err();
903 assert!(
905 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
906 "expected InvalidArgument(non-null), got {err:?}"
907 );
908 }
909
910 #[test]
911 fn from_foreign_ptr_rejects_empty_shape() {
912 let mut dummy: u8 = 0;
913 let err = unsafe {
914 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
915 }
916 .unwrap_err();
917 assert!(
918 matches!(err, crate::error::Error::InvalidSize(0)),
919 "expected InvalidSize(0) for empty shape, got {err:?}"
920 );
921 }
922
923 #[test]
924 fn from_foreign_ptr_rejects_overflow_shape() {
925 let mut dummy: u8 = 0;
926 let huge = [usize::MAX / 2 + 1, 2];
927 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
928 .unwrap_err();
929 assert!(
930 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
931 "expected InvalidArgument(overflow), got {err:?}"
932 );
933 }
934
935 #[test]
936 fn from_foreign_ptr_u8_dtype_dispatch() {
937 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
940 let ptr = buf.as_mut_ptr();
941 let owner: crate::ForeignOwner = Box::new(buf);
942 let t = unsafe {
943 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
944 }
945 .unwrap();
946 assert_eq!(t.dtype(), DType::U8);
947 assert_eq!(t.shape(), &[4]);
948 let m = t.as_u8().unwrap().map().unwrap();
949 use crate::TensorMapTrait;
950 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
951 }
952
953 #[test]
954 fn downcast_ref() {
955 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
956 let dyn_t: TensorDyn = t.into();
957 assert!(dyn_t.as_u8().is_some());
958 assert!(dyn_t.as_i8().is_none());
959 }
960
961 #[test]
962 fn downcast_into() {
963 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
964 let dyn_t: TensorDyn = t.into();
965 let back = dyn_t.into_u8().unwrap();
966 assert_eq!(back.shape(), &[10]);
967 }
968
969 #[test]
970 fn image_accessors() {
971 let t = Tensor::<u8>::image(640, 480, PixelFormat::Rgba, None).unwrap();
972 let dyn_t: TensorDyn = t.into();
973 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
974 assert_eq!(dyn_t.width(), Some(640));
975 assert_eq!(dyn_t.height(), Some(480));
976 assert!(!dyn_t.is_multiplane());
977 }
978
979 #[test]
980 fn image_constructor() {
981 let dyn_t = TensorDyn::image(640, 480, PixelFormat::Rgb, DType::U8, None).unwrap();
982 assert_eq!(dyn_t.dtype(), DType::U8);
983 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
984 assert_eq!(dyn_t.width(), Some(640));
985 }
986
987 #[test]
988 fn image_constructor_i8() {
989 let dyn_t = TensorDyn::image(640, 480, PixelFormat::Rgb, DType::I8, None).unwrap();
990 assert_eq!(dyn_t.dtype(), DType::I8);
991 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
992 }
993
994 #[test]
995 fn set_format_packed() {
996 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
997 assert_eq!(t.format(), None);
998 t.set_format(PixelFormat::Rgb).unwrap();
999 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1000 assert_eq!(t.width(), Some(640));
1001 assert_eq!(t.height(), Some(480));
1002 }
1003
1004 #[test]
1005 fn set_format_planar() {
1006 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
1007 t.set_format(PixelFormat::PlanarRgb).unwrap();
1008 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
1009 assert_eq!(t.width(), Some(640));
1010 assert_eq!(t.height(), Some(480));
1011 }
1012
1013 #[test]
1014 fn set_format_rejects_wrong_shape() {
1015 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1016 assert!(t.set_format(PixelFormat::Rgb).is_err());
1017 }
1018
1019 #[test]
1020 fn with_format_builder() {
1021 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1022 .unwrap()
1023 .with_format(PixelFormat::Rgba)
1024 .unwrap();
1025 assert_eq!(t.format(), Some(PixelFormat::Rgba));
1026 assert_eq!(t.width(), Some(640));
1027 assert_eq!(t.height(), Some(480));
1028 }
1029
1030 #[cfg(target_os = "linux")]
1031 #[test]
1032 fn dmabuf_clone_mem_tensor_fails() {
1033 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1034 assert_eq!(t.memory(), TensorMemory::Mem);
1035 assert!(t.dmabuf_clone().is_err());
1036 }
1037
1038 #[cfg(target_os = "linux")]
1039 #[test]
1040 fn dmabuf_mem_tensor_fails() {
1041 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1042 assert!(t.dmabuf().is_err());
1043 }
1044
1045 #[test]
1046 fn set_format_semi_planar_nv12() {
1047 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1049 t.set_format(PixelFormat::Nv12).unwrap();
1050 assert_eq!(t.format(), Some(PixelFormat::Nv12));
1051 assert_eq!(t.width(), Some(640));
1052 assert_eq!(t.height(), Some(480));
1053 }
1054
1055 #[test]
1056 fn set_format_semi_planar_nv16() {
1057 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1059 t.set_format(PixelFormat::Nv16).unwrap();
1060 assert_eq!(t.format(), Some(PixelFormat::Nv16));
1061 assert_eq!(t.width(), Some(640));
1062 assert_eq!(t.height(), Some(480));
1063 }
1064
1065 #[test]
1066 fn with_format_rejects_wrong_shape() {
1067 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1068 .unwrap()
1069 .with_format(PixelFormat::Rgb);
1070 assert!(result.is_err());
1071 }
1072
1073 #[test]
1074 fn set_format_preserved_after_rejection() {
1075 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1076 t.set_format(PixelFormat::Rgb).unwrap();
1077 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1078
1079 assert!(t.set_format(PixelFormat::Rgba).is_err());
1081
1082 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1084 }
1085
1086 #[test]
1087 fn set_format_idempotent() {
1088 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1089 t.set_format(PixelFormat::Rgb).unwrap();
1090 t.set_format(PixelFormat::Rgb).unwrap();
1091 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1092 assert_eq!(t.width(), Some(640));
1093 assert_eq!(t.height(), Some(480));
1094 }
1095
1096 #[test]
1099 fn set_row_stride_valid() {
1100 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1102 t.set_row_stride(512).unwrap();
1103 assert_eq!(t.row_stride(), Some(512));
1104 assert_eq!(t.effective_row_stride(), Some(512));
1105 }
1106
1107 #[test]
1108 fn set_row_stride_equals_min() {
1109 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1111 t.set_row_stride(300).unwrap();
1112 assert_eq!(t.row_stride(), Some(300));
1113 }
1114
1115 #[test]
1116 fn set_row_stride_too_small() {
1117 let mut t = TensorDyn::image(64, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1122 assert!(t.set_row_stride(200).is_err());
1123 assert_eq!(t.row_stride(), None);
1124 }
1125
1126 #[test]
1127 fn set_row_stride_zero() {
1128 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1129 assert!(t.set_row_stride(0).is_err());
1130 }
1131
1132 #[test]
1133 fn set_row_stride_requires_format() {
1134 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1135 assert!(t.set_row_stride(2048).is_err());
1136 }
1137
1138 #[test]
1139 fn effective_row_stride_without_stride() {
1140 let t = TensorDyn::image(64, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1145 assert_eq!(t.row_stride(), None);
1146 assert_eq!(t.effective_row_stride(), Some(192)); }
1148
1149 #[test]
1150 fn effective_row_stride_padded_packed_dma() {
1151 let t = match TensorDyn::image(
1157 100,
1158 100,
1159 PixelFormat::Rgb,
1160 DType::U8,
1161 Some(TensorMemory::Dma),
1162 ) {
1163 Ok(t) if t.memory() == TensorMemory::Dma => t,
1164 _ => return,
1165 };
1166 assert_eq!(t.row_stride(), Some(320));
1167 assert_eq!(t.effective_row_stride(), Some(320));
1168 }
1169
1170 #[test]
1171 fn effective_row_stride_no_format() {
1172 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1173 assert_eq!(t.effective_row_stride(), None);
1174 }
1175
1176 #[test]
1177 fn with_row_stride_builder() {
1178 let t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None)
1179 .unwrap()
1180 .with_row_stride(512)
1181 .unwrap();
1182 assert_eq!(t.row_stride(), Some(512));
1183 assert_eq!(t.effective_row_stride(), Some(512));
1184 }
1185
1186 #[test]
1187 fn with_row_stride_rejects_small() {
1188 let result = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None)
1189 .unwrap()
1190 .with_row_stride(200);
1191 assert!(result.is_err());
1192 }
1193
1194 #[test]
1195 fn set_format_clears_row_stride() {
1196 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1197 t.set_format(PixelFormat::Rgb).unwrap();
1198 t.set_row_stride(2048).unwrap();
1199 assert_eq!(t.row_stride(), Some(2048));
1200
1201 let _ = t.set_format(PixelFormat::Bgra);
1203 assert_eq!(t.row_stride(), Some(2048));
1204
1205 t.set_format(PixelFormat::Rgb).unwrap();
1207 assert_eq!(t.row_stride(), Some(2048));
1208
1209 t.reshape(&[480 * 640 * 3]).unwrap();
1211 assert_eq!(t.row_stride(), None);
1212 assert_eq!(t.format(), None);
1213 }
1214
1215 #[test]
1216 fn set_format_different_compatible_clears_stride() {
1217 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1220 t.set_format(PixelFormat::Rgba).unwrap();
1221 t.set_row_stride(4096).unwrap();
1222 assert_eq!(t.row_stride(), Some(4096));
1223
1224 t.set_format(PixelFormat::Bgra).unwrap();
1226 assert_eq!(t.format(), Some(PixelFormat::Bgra));
1227 assert_eq!(t.row_stride(), None);
1228 }
1229
1230 #[test]
1231 fn set_format_same_preserves_stride() {
1232 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1233 t.set_row_stride(512).unwrap();
1234 t.set_format(PixelFormat::Rgb).unwrap();
1236 assert_eq!(t.row_stride(), Some(512));
1237 }
1238
1239 #[test]
1240 fn effective_row_stride_planar() {
1241 let t = TensorDyn::image(640, 480, PixelFormat::PlanarRgb, DType::U8, None).unwrap();
1242 assert_eq!(t.effective_row_stride(), Some(640)); }
1244
1245 #[test]
1246 fn effective_row_stride_nv12() {
1247 let t = TensorDyn::image(640, 480, PixelFormat::Nv12, DType::U8, None).unwrap();
1248 assert_eq!(t.effective_row_stride(), Some(640)); }
1250
1251 #[test]
1252 fn map_rejects_strided_tensor() {
1253 let mut t =
1254 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1255 assert!(t.map().is_ok());
1257 t.set_row_stride(512).unwrap();
1259 let err = t.map();
1260 assert!(err.is_err());
1261 }
1262
1263 #[test]
1266 fn plane_offset_default_none() {
1267 let t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1268 assert_eq!(t.plane_offset(), None);
1269 }
1270
1271 #[test]
1272 fn set_plane_offset_basic() {
1273 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1274 t.set_plane_offset(4096);
1275 assert_eq!(t.plane_offset(), Some(4096));
1276 }
1277
1278 #[test]
1279 fn set_plane_offset_zero() {
1280 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1281 t.set_plane_offset(0);
1282 assert_eq!(t.plane_offset(), Some(0));
1283 }
1284
1285 #[test]
1286 fn set_plane_offset_no_format() {
1287 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1289 t.set_plane_offset(4096);
1290 assert_eq!(t.plane_offset(), Some(4096));
1291 }
1292
1293 #[test]
1294 fn set_format_clears_plane_offset() {
1295 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1296 t.set_format(PixelFormat::Rgb).unwrap();
1297 t.set_plane_offset(4096);
1298 assert_eq!(t.plane_offset(), Some(4096));
1299
1300 t.set_format(PixelFormat::Rgb).unwrap();
1302 assert_eq!(t.plane_offset(), Some(4096));
1303
1304 t.reshape(&[480 * 640 * 3]).unwrap();
1306 assert_eq!(t.plane_offset(), None);
1307 assert_eq!(t.format(), None);
1308 }
1309
1310 #[test]
1311 fn map_rejects_out_of_bounds_offset() {
1312 let mut t =
1313 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1314 assert!(t.map().is_ok());
1316 t.set_plane_offset(4096);
1319 assert!(t.map().is_err());
1320 }
1321
1322 #[test]
1323 fn mem_subview_in_bounds_maps_at_offset() {
1324 let parent =
1327 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1328 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
1330 assert_eq!(view.plane_offset(), Some(4096));
1331 assert!(view.map().is_ok());
1332 }
1333
1334 #[test]
1335 fn dyn_batch_dispatches_every_dtype() {
1336 use DType::*;
1340 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
1341 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
1342 let view = parent.batch(1).unwrap();
1343 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
1344 assert_eq!(view.shape(), &[4], "{dt:?}");
1345 }
1346 }
1347
1348 #[test]
1349 fn map_accepts_zero_offset_tensor() {
1350 let mut t =
1351 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1352 t.set_plane_offset(0);
1353 assert!(t.map().is_ok());
1355 }
1356
1357 #[test]
1358 fn dyn_configure_image_nv12() {
1359 let mut t = TensorDyn::image(640, 480, PixelFormat::Rgb, DType::U8, None).unwrap();
1360 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
1361 assert_eq!(t.format(), Some(PixelFormat::Nv12));
1362 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
1363 }
1364
1365 #[test]
1366 fn tensordyn_colorimetry_roundtrip() {
1367 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
1368 let mut t = TensorDyn::image(1280, 720, PixelFormat::Nv12, DType::U8, None).unwrap();
1369 assert_eq!(t.colorimetry(), None);
1370 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
1371 t.set_colorimetry(Some(c));
1372 assert_eq!(t.colorimetry(), Some(c));
1373 }
1374
1375 #[test]
1376 fn from_planes_propagates_plane_offset() {
1377 let mut luma =
1378 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
1379 luma.set_plane_offset(4096);
1380 let chroma =
1381 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
1382 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
1383 assert_eq!(combined.plane_offset(), Some(4096));
1384 }
1385
1386 #[test]
1387 fn cuda_passthrough_none_for_mem_tensor() {
1388 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
1391 .unwrap()
1392 .into();
1393 assert!(t.cuda().is_none());
1394 assert!(t.cuda_map().is_none());
1395 }
1396}