mnn/
tensor.rs

1use crate::prelude::*;
2use core::marker::PhantomData;
3pub mod list;
4mod tensor_any;
5mod tensor_ref;
6pub use tensor_any::AnyTensorRef;
7pub use tensor_ref::{TensorRef, from_raw_parts, from_raw_parts_mut};
8
9use mnn_sys::HalideType;
10
11/// A tensor that can be owned by the host or device
12pub type TensorOwned<H, M> = Tensor<Owned<H>, M>;
13/// A tensor that can be borrowed from the host or device
14pub type TensorView<'t, H, M> = Tensor<View<&'t H>, M>;
15/// A tensor that can be borrowed mutably from the host or device
16pub type TensorViewMut<'t, H, M> = Tensor<View<&'t mut H>, M>;
17
18mod seal {
19    pub trait Sealed {}
20}
21impl seal::Sealed for Host {}
22impl seal::Sealed for Device {}
23impl<T> seal::Sealed for View<T> {}
24impl<T> seal::Sealed for Owned<T> {}
25
26/// A trait to represent the type of a tensor
27pub trait TensorType: seal::Sealed {
28    /// The type of the tensor data
29    type H;
30    /// Check if the tensor is owned
31    fn owned() -> bool;
32    /// Check if the tensor is borrowed
33    fn borrowed() -> bool {
34        !Self::owned()
35    }
36}
37
38/// A trait to represent a mutable tensor type
39pub trait MutableTensorType: TensorType + seal::Sealed {}
40
41impl<H> MutableTensorType for Owned<H> {}
42impl<H> MutableTensorType for View<&mut H> {}
43
44impl<H> TensorType for Owned<H> {
45    type H = H;
46    fn owned() -> bool {
47        true
48    }
49}
50
51impl<H> TensorType for View<&H> {
52    type H = H;
53    fn owned() -> bool {
54        false
55    }
56}
57
58impl<H> TensorType for View<&mut H> {
59    type H = H;
60    fn owned() -> bool {
61        false
62    }
63}
64
65/// A trait to represent the device type of a tensor
66pub trait TensorMachine: seal::Sealed {
67    /// Check if the tensor is owned by the device
68    fn device() -> bool;
69    /// Check if the tensor is owned by the host
70    fn host() -> bool {
71        !Self::device()
72    }
73}
74
75impl TensorMachine for Host {
76    fn device() -> bool {
77        false
78    }
79}
80
81impl TensorMachine for Device {
82    fn device() -> bool {
83        true
84    }
85}
86
87/// A tensor that is owned by the cpu / host platform
88#[non_exhaustive]
89#[derive(Debug)]
90pub struct Host {}
91/// A tensor that is owned by the device / gpu platform
92#[non_exhaustive]
93#[derive(Debug)]
94pub struct Device {}
95
96/// A reference to a any tensor
97#[repr(transparent)]
98#[derive(Debug)]
99pub struct View<T> {
100    pub(crate) __marker: PhantomData<[T]>,
101}
102/// A tensor that is owned by the host / device platform
103#[repr(transparent)]
104#[derive(Debug)]
105pub struct Owned<T> {
106    pub(crate) __marker: PhantomData<T>,
107}
108
109/// A generic tensor that can of host / device / owned / borrowed
110///
111/// Every `Tensor` owns the C++ `MNN::Tensor` object it wraps and destroys it on drop. The
112/// `S` parameter only describes who owns the *data*: [`Owned`] tensors allocate it, the
113/// [`View`] tensors created by [`Tensor::borrowed`] / [`Tensor::borrowed_mut`] borrow a user
114/// slice (which MNN never frees). To borrow a tensor itself use `&TensorRef` /
115/// `&mut TensorRef`, the [`Deref`](core::ops::Deref) target of `Tensor`.
116#[derive(Debug)]
117#[repr(transparent)]
118pub struct Tensor<S, M, A = <S as TensorType>::H>
119where
120    A: HalideType,
121    M: TensorMachine,
122    S: TensorType<H = A>,
123{
124    pub(crate) tensor: *mut mnn_sys::Tensor,
125    __marker: PhantomData<(S, M, A)>,
126}
127
128impl<S, M, A> Drop for Tensor<S, M, A>
129where
130    A: HalideType,
131    S: TensorType<H = A>,
132    M: TensorMachine,
133{
134    fn drop(&mut self) {
135        unsafe {
136            mnn_sys::Tensor_destroy(self.tensor);
137        }
138    }
139}
140
141/// The type of the tensor dimension
142/// If you are manually specifying the shapes then this doesn't really matter
143/// N -> Batch size
144/// C -> Channel
145/// H -> Height
146/// W -> Width
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
148pub enum DimensionType {
149    /// Caffe style dimensions or NCHW
150    #[default]
151    Caffe,
152    /// Caffe style dimensions with channel packed in 4 bytes or NC4HW4
153    CaffeC4,
154    /// Tensorflow style dimensions or NHWC
155    TensorFlow,
156}
157
158impl DimensionType {
159    /// Tensorflow style dimensions or NHWC
160    pub const NHWC: Self = Self::TensorFlow;
161    /// Caffe style dimensions or NCHW
162    pub const NCHW: Self = Self::Caffe;
163    /// Caffe style dimensions with channel packed in 4 bytes or NC4HW4
164    pub const NC4HW4: Self = Self::CaffeC4;
165    pub(crate) fn to_mnn_sys(self) -> mnn_sys::DimensionType {
166        match self {
167            DimensionType::Caffe => mnn_sys::DimensionType::CAFFE,
168            DimensionType::CaffeC4 => mnn_sys::DimensionType::CAFFE_C4,
169            DimensionType::TensorFlow => mnn_sys::DimensionType::TENSORFLOW,
170        }
171    }
172}
173
174impl From<mnn_sys::DimensionType> for DimensionType {
175    fn from(dm: mnn_sys::DimensionType) -> Self {
176        match dm {
177            mnn_sys::DimensionType::CAFFE => DimensionType::Caffe,
178            mnn_sys::DimensionType::CAFFE_C4 => DimensionType::CaffeC4,
179            mnn_sys::DimensionType::TENSORFLOW => DimensionType::TensorFlow,
180        }
181    }
182}
183
184impl<H, T, M> Tensor<T, M>
185where
186    T: TensorType<H = H>,
187    H: HalideType,
188    M: TensorMachine,
189{
190    /// This function constructs a Tensor type from a raw pointer
191    ///
192    /// The returned tensor takes ownership of the `MNN::Tensor` object and destroys it on
193    /// drop. Use [`from_raw_parts`] / [`from_raw_parts_mut`] to borrow a tensor instead.
194    ///# Safety
195    /// Since this constructs a Tensor from a raw pointer we have no way to guarantee that it's a
196    /// valid tensor, that nothing else owns it, or its lifetime
197    pub unsafe fn from_ptr(tensor: *mut mnn_sys::Tensor) -> Self {
198        assert!(!tensor.is_null());
199        Self {
200            tensor,
201            __marker: PhantomData,
202        }
203    }
204
205    /// Print the tensor
206    pub fn print(&self) {
207        unsafe {
208            mnn_sys::Tensor_print(self.tensor);
209        }
210    }
211
212    /// DO not use this function directly
213    /// # Safety
214    /// This is just provided as a 1:1 compat mostly for possible later use
215    pub unsafe fn halide_buffer(&self) -> *const mnn_sys::halide_buffer_t {
216        unsafe { mnn_sys::Tensor_buffer(self.tensor) }
217    }
218
219    /// Do not use this function directly
220    /// # Safety
221    /// This is just provided as a 1:1 compat mostly for possible later use
222    pub unsafe fn halide_buffer_mut(&self) -> *mut mnn_sys::halide_buffer_t {
223        unsafe { mnn_sys::Tensor_buffer_mut(self.tensor) }
224    }
225
226    /// Get the data type of the tensor
227    pub fn get_type(&self) -> mnn_sys::halide_type_t {
228        unsafe { mnn_sys::Tensor_getType(self.tensor) }
229    }
230
231    /// # Safety
232    ///
233    /// Type erase the tensor
234    pub unsafe fn as_any_tensor(&self) -> &AnyTensorRef {
235        unsafe { AnyTensorRef::from_ptr(self.tensor) }
236    }
237
238    /// # Safety
239    ///
240    /// Type erase the tensor mutably
241    pub unsafe fn as_any_tensor_mut(&mut self) -> &mut AnyTensorRef {
242        unsafe { AnyTensorRef::from_mut_ptr(self.tensor) }
243    }
244}
245
246impl<H> Tensor<Owned<H>, Host>
247where
248    H: HalideType + Copy,
249{
250    /// Create a new tensor with the specified shape and dimension type and fill it with the data from the provided vector
251    ///
252    /// # Panics
253    /// Panics if the shape is dynamic or the number of elements MNN allocates for the shape
254    /// (which includes channel padding for [`DimensionType::CaffeC4`]) differs from `data.len()`.
255    pub fn from_shape_data(shape: impl AsTensorShape, dm_type: DimensionType, data: &[H]) -> Self {
256        let shape = shape.as_tensor_shape();
257        let size = shape
258            .tensor_size()
259            .unwrap_or_else(|| panic!("Tensor shape {shape:?} is dynamic or too large"));
260        assert_eq!(
261            size,
262            data.len(),
263            "Data length ({}) does not match tensor shape ({:?})",
264            data.len(),
265            shape
266        );
267        let mut tensor = Self::new(shape, dm_type);
268        let host = tensor.host_mut();
269        assert_eq!(
270            host.len(),
271            data.len(),
272            "Tensor layout {dm_type:?} needs {} elements for shape {shape:?} (channel padding) but {} were provided",
273            host.len(),
274            data.len()
275        );
276        host.copy_from_slice(data);
277        tensor
278    }
279}
280
281impl<H, M> Tensor<Owned<H>, M>
282where
283    H: HalideType,
284    M: TensorMachine,
285{
286    /// Create a new tensor with the specified shape and dimension type this allocates the tensor
287    /// internally and the data is "owned" by the tensor itself.
288    ///
289    /// # Panics
290    /// Host tensors are allocated immediately, so their shape must not be dynamic (contain
291    /// `-1`). Device tensors are only described, not allocated, so any shape is accepted.
292    pub fn new(shape: impl AsTensorShape, dm_type: DimensionType) -> Self {
293        let shape = shape.as_tensor_shape();
294        if M::host() {
295            assert!(
296                shape.tensor_size().is_some(),
297                "Cannot allocate a host tensor with a dynamic or oversized shape {shape:?}"
298            );
299        }
300        let tensor = unsafe {
301            if M::device() {
302                mnn_sys::Tensor_createDevice(
303                    shape.shape.as_ptr(),
304                    shape.size,
305                    mnn_sys::halide_type_of::<H>(),
306                    dm_type.to_mnn_sys(),
307                )
308            } else {
309                mnn_sys::Tensor_createWith(
310                    shape.shape.as_ptr(),
311                    shape.size,
312                    mnn_sys::halide_type_of::<H>(),
313                    core::ptr::null_mut(),
314                    dm_type.to_mnn_sys(),
315                )
316            }
317        };
318        debug_assert!(!tensor.is_null());
319        Self {
320            tensor,
321            __marker: PhantomData,
322        }
323    }
324}
325
326impl<H> Clone for Tensor<Owned<H>, Host>
327where
328    H: HalideType + Copy,
329{
330    fn clone(&self) -> Tensor<Owned<H>, Host> {
331        let data = self.host();
332        let shape = self.shape();
333        let mut out = Tensor::new(shape, self.get_dimension_type());
334        out.host_mut().copy_from_slice(data);
335        // Cloning / deepCopy  is not supported by mnn currently https://github.com/alibaba/MNN/blob/6b1db4c2eacf8193e6b7a1841acb6e31a1938e82/include/MNN/Tensor.hpp#L131
336
337        out
338    }
339}
340
341impl<T, H> PartialEq for Tensor<T, Host>
342where
343    T: TensorType<H = H>,
344    H: HalideType + PartialEq,
345{
346    fn eq(&self, other: &Self) -> bool {
347        if self.size() != other.size() {
348            return false;
349        }
350        let self_data = self.host();
351        let other_data = other.host();
352        self_data == other_data
353    }
354}
355
356impl<T, H> Eq for Tensor<T, Host>
357where
358    T: TensorType<H = H>,
359    H: HalideType + Eq,
360{
361}
362
363/// A tensor shape
364#[derive(Clone, Copy)]
365#[repr(C)]
366pub struct TensorShape {
367    pub(crate) shape: [i32; 4],
368    pub(crate) size: usize,
369}
370
371impl TensorShape {
372    /// Number of elements described by the shape (the product of all the extents).
373    ///
374    /// Returns `None` if the shape is dynamic (contains a negative extent such as `-1`) or
375    /// if the product overflows.
376    pub fn tensor_size(&self) -> Option<usize> {
377        self.shape[..self.size].iter().try_fold(1usize, |acc, &x| {
378            let x = usize::try_from(x).ok()?;
379            acc.checked_mul(x)
380        })
381    }
382}
383
384impl From<mnn_sys::TensorShape> for TensorShape {
385    fn from(value: mnn_sys::TensorShape) -> Self {
386        Self {
387            shape: value.shape,
388            size: value.size,
389        }
390    }
391}
392
393impl From<TensorShape> for mnn_sys::TensorShape {
394    fn from(value: TensorShape) -> Self {
395        Self {
396            shape: value.shape,
397            size: value.size,
398        }
399    }
400}
401
402impl core::ops::Deref for TensorShape {
403    type Target = [i32];
404
405    fn deref(&self) -> &Self::Target {
406        &self.shape[..self.size]
407    }
408}
409
410impl core::ops::Index<usize> for TensorShape {
411    type Output = i32;
412
413    fn index(&self, index: usize) -> &Self::Output {
414        &self.shape[..self.size][index]
415    }
416}
417
418impl core::ops::IndexMut<usize> for TensorShape {
419    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
420        &mut self.shape[..self.size][index]
421    }
422}
423
424impl core::ops::DerefMut for TensorShape {
425    fn deref_mut(&mut self) -> &mut Self::Target {
426        &mut self.shape[..self.size]
427    }
428}
429
430impl core::fmt::Debug for TensorShape {
431    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
432        f.debug_tuple("TensorShape")
433            .field(&&self.shape[..self.size])
434            .finish()
435    }
436}
437
438/// A trait to convert any array-like type to a tensor shape
439pub trait AsTensorShape {
440    /// Convert the array-like type to a tensor shape
441    fn as_tensor_shape(&self) -> TensorShape;
442}
443
444impl<T: AsRef<[i32]>> AsTensorShape for T {
445    fn as_tensor_shape(&self) -> TensorShape {
446        let this = self.as_ref();
447        let size = std::cmp::min(this.len(), 4);
448        let mut shape = [1; 4];
449        shape[..size].copy_from_slice(&this[..size]);
450        TensorShape { shape, size }
451    }
452}
453
454impl AsTensorShape for TensorShape {
455    fn as_tensor_shape(&self) -> TensorShape {
456        *self
457    }
458}
459
460#[cfg(test)]
461mod as_tensor_shape_tests {
462    use super::AsTensorShape;
463    macro_rules! shape_test {
464        ($t:ty, $kind: expr, $value: expr) => {
465            eprintln!("Testing {} with {} shape", stringify!($t), $kind);
466            $value.as_tensor_shape();
467        };
468    }
469    #[test]
470    fn as_tensor_shape_test_vec() {
471        shape_test!(Vec<i32>, "small", vec![1, 2, 3]);
472        shape_test!(Vec<i32>, "large", vec![12, 23, 34, 45, 67]);
473    }
474    #[test]
475    fn as_tensor_shape_test_array() {
476        shape_test!([i32; 3], "small", [1, 2, 3]);
477        shape_test!([i32; 5], "large", [12, 23, 34, 45, 67]);
478    }
479    #[test]
480    fn as_tensor_shape_test_ref() {
481        shape_test!(&[i32], "small", &[1, 2, 3]);
482        shape_test!(&[i32], "large", &[12, 23, 34, 45, 67]);
483    }
484
485    #[test]
486    fn tensor_size_handles_dynamic_and_overflow() {
487        use super::TensorShape;
488        assert_eq!([1, 2, 3].as_tensor_shape().tensor_size(), Some(6));
489        assert_eq!([1, -1, 3].as_tensor_shape().tensor_size(), None);
490        assert_eq!([i32::MAX; 4].as_tensor_shape().tensor_size(), None);
491        // Debug formatting must not panic on dynamic shapes.
492        let dynamic = TensorShape {
493            shape: [1, 3, -1, 0],
494            size: 3,
495        };
496        assert_eq!(format!("{dynamic:?}"), "TensorShape([1, 3, -1])");
497    }
498}
499
500#[cfg(test)]
501mod tensor_tests {
502    #[test]
503    #[should_panic]
504    fn unsafe_nullptr_tensor() {
505        use super::*;
506        unsafe {
507            super::Tensor::<Owned<i32>, Host>::from_ptr(core::ptr::null_mut());
508        }
509    }
510
511    /// Verifies that `try_host()` is NOT callable on Device tensors.
512    /// Previously this was a soundness hole — calling try_host() on a Device
513    /// tensor would dereference a null pointer (segfault). The fix restricts
514    /// try_host/host to `Tensor<T, Host>` only.
515    #[test]
516    fn try_host_not_available_on_device_tensor() {
517        use super::*;
518        let device_tensor = Tensor::<Owned<f32>, Device>::new([1, 2, 3], DimensionType::Caffe);
519        // try_host() should only exist on Host tensors, not Device tensors.
520        // Uncommenting the next line must fail to compile:
521        // let _ = device_tensor.try_host();
522        drop(device_tensor);
523    }
524}
525
526impl<'a, A> Tensor<View<&'a A>, Host, A>
527where
528    A: HalideType,
529{
530    /// Create a tensor that borrows `input` as its storage.
531    ///
532    /// # Panics
533    /// Panics if the shape is dynamic or if the number of elements MNN needs for `shape` in
534    /// the `dm_type` layout differs from `input.len()`. Note that [`DimensionType::CaffeC4`]
535    /// pads the channel dimension to a multiple of 4, so the slice must already be laid out
536    /// (and sized) accordingly.
537    pub fn borrowed(shape: impl AsTensorShape, dm_type: DimensionType, input: &'a [A]) -> Self {
538        let shape = shape.as_tensor_shape();
539        let size = shape
540            .tensor_size()
541            .unwrap_or_else(|| panic!("Tensor shape {shape:?} is dynamic or too large"));
542        assert_eq!(
543            size,
544            input.len(),
545            "Input data length ({}) does not match the tensor shape ({:?})",
546            input.len(),
547            shape
548        );
549        let tensor = unsafe {
550            mnn_sys::Tensor_createWith(
551                shape.shape.as_ptr(),
552                shape.size,
553                mnn_sys::halide_type_of::<A>(),
554                input.as_ptr().cast_mut().cast(),
555                dm_type.to_mnn_sys(),
556            )
557        };
558        assert!(!tensor.is_null());
559        // MNN never frees memory it did not allocate, so destroying this wrapper on drop
560        // leaves `input` untouched.
561        let output = Self {
562            tensor,
563            __marker: PhantomData,
564        };
565        // `output` is dropped (and the wrapper destroyed) if this panics.
566        assert_eq!(
567            output.element_size(),
568            input.len(),
569            "Tensor layout {dm_type:?} needs {} elements for shape {shape:?} (channel padding) but the slice has {}",
570            output.element_size(),
571            input.len()
572        );
573        output
574    }
575}
576
577impl<'a, A> Tensor<View<&'a mut A>, Host, A>
578where
579    A: HalideType,
580{
581    /// Create a tensor that mutably borrows `input` as its storage.
582    ///
583    /// # Panics
584    /// Panics if the shape is dynamic or if the number of elements MNN needs for `shape` in
585    /// the `dm_type` layout differs from `input.len()`. Note that [`DimensionType::CaffeC4`]
586    /// pads the channel dimension to a multiple of 4, so the slice must already be laid out
587    /// (and sized) accordingly.
588    pub fn borrowed_mut(
589        shape: impl AsTensorShape,
590        dm_type: DimensionType,
591        input: &'a mut [A],
592    ) -> Self {
593        let shape = shape.as_tensor_shape();
594        let size = shape
595            .tensor_size()
596            .unwrap_or_else(|| panic!("Tensor shape {shape:?} is dynamic or too large"));
597        assert_eq!(
598            size,
599            input.len(),
600            "Input data length ({}) does not match the tensor shape ({:?})",
601            input.len(),
602            shape
603        );
604        let len = input.len();
605        let tensor = unsafe {
606            mnn_sys::Tensor_createWith(
607                shape.shape.as_ptr(),
608                shape.size,
609                mnn_sys::halide_type_of::<A>(),
610                input.as_mut_ptr().cast(),
611                dm_type.to_mnn_sys(),
612            )
613        };
614        assert!(!tensor.is_null());
615        // MNN never frees memory it did not allocate, so destroying this wrapper on drop
616        // leaves `input` untouched.
617        let output = Self {
618            tensor,
619            __marker: PhantomData,
620        };
621        // `output` is dropped (and the wrapper destroyed) if this panics.
622        assert_eq!(
623            output.element_size(),
624            len,
625            "Tensor layout {dm_type:?} needs {} elements for shape {shape:?} (channel padding) but the slice has {}",
626            output.element_size(),
627            len
628        );
629        output
630    }
631}
632
633#[test]
634fn test_tensor_borrowed() {
635    let shape = [1, 2, 3];
636    let data = vec![1, 2, 3, 4, 5, 6];
637    let tensor = Tensor::<View<&i32>, Host>::borrowed(shape, DimensionType::Caffe, &data);
638    assert_eq!(tensor.shape().as_ref(), shape);
639    assert_eq!(tensor.host(), data.as_slice());
640}
641#[test]
642fn test_tensor_borrow_mut() {
643    let shape = [1, 2, 3];
644    let mut data = vec![1, 2, 3, 4, 5, 6];
645    let mut tensor =
646        Tensor::<View<&mut i32>, Host>::borrowed_mut(shape, DimensionType::Caffe, &mut data);
647    tensor.host_mut().fill(1);
648    drop(tensor);
649    assert_eq!(data, &[1, 1, 1, 1, 1, 1]);
650}
651
652#[test]
653#[should_panic]
654fn test_tensor_invalid_size() {
655    let shape = [300, 400, 500];
656    let date = vec![0; 100];
657    Tensor::<View<&i32>, Host>::borrowed(shape, DimensionType::default(), &date);
658}
659
660#[test]
661#[should_panic(expected = "channel padding")]
662fn test_tensor_borrowed_rejects_unpadded_c4() {
663    // NC4HW4 pads the channel dimension (3) up to 4, so MNN would read 16 elements
664    // from a 12 element slice.
665    let data = vec![1f32; 12];
666    Tensor::<View<&f32>, Host>::borrowed([1, 3, 2, 2], DimensionType::CaffeC4, &data);
667}
668
669#[test]
670#[should_panic(expected = "channel padding")]
671fn test_tensor_borrowed_mut_rejects_unpadded_c4() {
672    let mut data = vec![1f32; 12];
673    Tensor::<View<&mut f32>, Host>::borrowed_mut([1, 3, 2, 2], DimensionType::CaffeC4, &mut data);
674}
675
676#[test]
677fn test_tensor_borrowed_padded_c4() {
678    // With a channel count that is a multiple of 4 the layouts agree.
679    let data = vec![1f32; 16];
680    let tensor = Tensor::<View<&f32>, Host>::borrowed([1, 4, 2, 2], DimensionType::CaffeC4, &data);
681    assert_eq!(tensor.host().len(), 16);
682}
683
684#[test]
685#[should_panic(expected = "channel padding")]
686fn test_from_shape_data_rejects_unpadded_c4() {
687    TensorOwned::<f32, Host>::from_shape_data([1, 3, 2, 2], DimensionType::CaffeC4, &[0f32; 12]);
688}
689
690#[test]
691#[should_panic(expected = "dynamic")]
692fn test_tensor_borrowed_rejects_dynamic_shape() {
693    let data = vec![0i32; 6];
694    Tensor::<View<&i32>, Host>::borrowed([1, -1, 3], DimensionType::Caffe, &data);
695}
696
697#[test]
698fn test_tensor_ref_borrows() {
699    let data = vec![1, 2, 3, 4, 5, 6];
700    let mut owned =
701        TensorOwned::<i32, Host>::from_shape_data([1, 2, 3], DimensionType::Caffe, &data);
702    {
703        let shared: &TensorRef<i32, Host> = &owned;
704        let again: &TensorRef<i32, Host> = shared;
705        assert_eq!(shared.host(), data.as_slice());
706        assert_eq!(again.host(), data.as_slice());
707    }
708    {
709        let exclusive: &mut TensorRef<i32, Host> = &mut owned;
710        exclusive.host_mut().fill(7);
711        let reborrowed: &mut TensorRef<i32, Host> = &mut *exclusive;
712        reborrowed.host_mut()[0] = 9;
713    }
714    // The owner is untouched by its borrows going out of scope.
715    assert_eq!(owned.host(), &[9, 7, 7, 7, 7, 7]);
716}
717
718#[test]
719fn test_borrowed_view_drop_keeps_user_data() {
720    let data = vec![1, 2, 3, 4, 5, 6];
721    for _ in 0..3 {
722        let view = Tensor::<View<&i32>, Host>::borrowed([1, 2, 3], DimensionType::Caffe, &data);
723        assert_eq!(view.host(), data.as_slice());
724        // Dropping the view destroys only the MNN wrapper, never `data`.
725    }
726    assert_eq!(data, &[1, 2, 3, 4, 5, 6]);
727}
728
729#[test]
730fn test_tensor_clone() {
731    let shape = [1, 2, 3];
732    let data = vec![1, 3, 5, 8, 1, 23];
733    let mut tensor = Tensor::new(shape, DimensionType::Caffe);
734    tensor.host_mut().copy_from_slice(&data);
735
736    let tensor_cloned = tensor.clone();
737    assert_eq!(tensor_cloned.host(), data);
738}
739
740#[test]
741fn test_tensor_creation() {
742    let shape = [4, 2, 8];
743    let data: &[u8; 64] = b"4756ff762c7a7909a1ec0015836296807602fb2d44e1cac15e057e22b2258040";
744
745    let tensor = {
746        let data = data.iter().copied().map(|v| v as f32).collect::<Vec<_>>();
747        TensorOwned::<f32, Host>::from_shape_data(shape, DimensionType::Caffe, &data)
748    };
749    let data = data.iter().copied().map(|v| v as f32).collect::<Vec<_>>();
750    assert_eq!(tensor.host(), data);
751}
mnn/tensor.rs

mnn/
tensor.rs
