rten_tensor/
contiguous.rs

1use std::ops::Deref;
2
3use crate::layout::FromShape;
4use crate::storage::{CowData, ViewData};
5use crate::{AsView, Layout, Storage, TensorBase};
6
7/// A tensor wrapper which guarantees that the tensor has a contiguous layout.
8///
9/// A contiguous layout means that the order of elements in memory matches the
10/// logical row-major ordering of elements with no gaps.
11#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
12pub struct Contiguous<T>(T);
13
14impl<T> Deref for Contiguous<T> {
15    type Target = T;
16
17    fn deref(&self) -> &T {
18        &self.0
19    }
20}
21
22impl<T> Contiguous<T> {
23    /// Extract the tensor from the wrapper.
24    pub fn into_inner(self) -> T {
25        self.0
26    }
27}
28
29impl<S: Storage, L: Layout> Contiguous<TensorBase<S, L>> {
30    /// Wrap a tensor if it is contiguous, or return `None` if the tensor has
31    /// a non-contiguous layout.
32    pub fn new(inner: TensorBase<S, L>) -> Option<Self> {
33        if inner.is_contiguous() {
34            Some(Self(inner))
35        } else {
36            None
37        }
38    }
39
40    /// Return the tensor's underlying data as a slice.
41    ///
42    /// Unlike [`TensorBase::data`] this returns a slice instead of an option
43    /// because the tensor is known to be contiguous.
44    pub fn data(&self) -> &[S::Elem] {
45        let len = self.0.len();
46        let ptr = self.0.data_ptr();
47
48        // Safety: Constructor verified that tensor is contiguous.
49        unsafe { std::slice::from_raw_parts(ptr, len) }
50    }
51
52    /// Return a contiguous view of this tensor.
53    pub fn view(&self) -> Contiguous<TensorBase<ViewData<'_, S::Elem>, L>>
54    where
55        TensorBase<S, L>: AsView<Elem = S::Elem, Layout = L>,
56    {
57        Contiguous(self.0.view())
58    }
59}
60
61impl<T, L: Clone + Layout> Contiguous<TensorBase<Vec<T>, L>> {
62    /// Wrap `inner` as a contiguous tensor.
63    ///
64    /// This is cheap if `inner` is already contiguous, otherwise the elements
65    /// are copied into a new buffer.
66    pub fn from_owned(mut inner: TensorBase<Vec<T>, L>) -> Self
67    where
68        L: FromShape,
69        T: Clone,
70    {
71        inner.make_contiguous();
72        Self(inner)
73    }
74
75    /// Extract the owned, contiguous data from this tensor.
76    pub fn into_data(self) -> Vec<T> {
77        self.0.into_non_contiguous_data()
78    }
79}
80
81impl<'a, T, L: Clone + Layout> Contiguous<TensorBase<CowData<'a, T>, L>> {
82    /// Extract the owned data from this tensor, if the data is owned.
83    pub fn into_data(self) -> Option<Vec<T>> {
84        self.0.into_non_contiguous_data()
85    }
86}
87
88impl<S: Storage, L: Layout> From<Contiguous<TensorBase<S, L>>> for TensorBase<S, L> {
89    fn from(val: Contiguous<TensorBase<S, L>>) -> Self {
90        val.0
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use crate::{AsView, Contiguous, Layout, NdTensor};
97
98    #[test]
99    fn test_contiguous() {
100        let tensor = NdTensor::<f32, 2>::zeros([3, 3]);
101        let wrapped = Contiguous::new(tensor);
102        assert!(wrapped.is_some());
103
104        let mut tensor: NdTensor<f32, 2> = wrapped.unwrap().into();
105        tensor.transpose();
106        let wrapped = Contiguous::new(tensor);
107        assert!(wrapped.is_none());
108    }
109
110    #[test]
111    fn test_contiguous_view() {
112        let tensor = NdTensor::<f32, 2>::zeros([3, 4]);
113        let wrapped = Contiguous::new(tensor).unwrap();
114        assert_eq!(wrapped.view().shape(), [3, 4]);
115    }
116}
rten_tensor/contiguous.rs

rten_tensor/
contiguous.rs
