rstsr_core/tensor/manuplication/
to_layout.rs

1use crate::prelude_dev::*;
2
3/* #region to_layout */
4
5pub fn change_layout_f<'a, R, T, B, D, D2>(
6    tensor: TensorAny<R, T, B, D>,
7    layout: Layout<D2>,
8) -> Result<TensorCow<'a, T, B, D2>>
9where
10    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
11    D: DimAPI,
12    D2: DimAPI,
13    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D2, D>,
14{
15    let shape = layout.shape();
16    rstsr_assert_eq!(tensor.size(), shape.shape_size(), InvalidLayout)?;
17    let same_layout = tensor.layout().to_dim::<IxD>()? == layout.to_dim::<IxD>()?;
18    let contig_c = tensor.c_contig() && layout.c_contig() && tensor.layout().offset() == layout.offset();
19    let contig_f = tensor.f_contig() && layout.f_contig() && tensor.layout().offset() == layout.offset();
20    let default_order = tensor.device().default_order();
21    let contig = match default_order {
22        RowMajor => contig_c,
23        ColMajor => contig_f,
24    };
25    if same_layout || contig {
26        // no data cloned
27        let (storage, _) = tensor.into_raw_parts();
28        let tensor = unsafe { TensorBase::new_unchecked(storage, layout) };
29        return Ok(tensor.into_cow());
30    } else {
31        // layout changed, or not c and f contiguous with same layout
32        // clone data by assign
33        let (storage_old, layout_old) = tensor.into_raw_parts();
34        let device = storage_old.device();
35        let (_, idx_max) = layout.bounds_index()?;
36        let mut storage_new = device.uninit_impl(idx_max)?;
37        device.assign_arbitary_uninit(storage_new.raw_mut(), &layout, storage_old.raw(), &layout_old)?;
38        let storage_new = unsafe { B::assume_init_impl(storage_new)? };
39        let tensor = unsafe { TensorBase::new_unchecked(storage_new, layout) };
40        return Ok(tensor.into_cow());
41    }
42}
43
44/// Convert tensor to the other layout.
45pub fn to_layout<R, T, D, B, D2>(tensor: &TensorAny<R, T, B, D>, layout: Layout<D2>) -> TensorCow<'_, T, B, D2>
46where
47    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
48    D: DimAPI,
49    D2: DimAPI,
50    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D2, D>,
51{
52    change_layout_f(tensor.view(), layout).rstsr_unwrap()
53}
54
55pub fn to_layout_f<R, T, D, B, D2>(
56    tensor: &TensorAny<R, T, B, D>,
57    layout: Layout<D2>,
58) -> Result<TensorCow<'_, T, B, D2>>
59where
60    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
61    D: DimAPI,
62    D2: DimAPI,
63    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D2, D>,
64{
65    change_layout_f(tensor.view(), layout)
66}
67
68pub fn into_layout_f<'a, R, T, B, D, D2>(tensor: TensorAny<R, T, B, D>, layout: Layout<D2>) -> Result<Tensor<T, B, D2>>
69where
70    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
71    D: DimAPI,
72    D2: DimAPI,
73    T: Clone,
74    B: DeviceAPI<T>
75        + DeviceRawAPI<MaybeUninit<T>>
76        + DeviceCreationAnyAPI<T>
77        + OpAssignArbitaryAPI<T, D2, D>
78        + OpAssignAPI<T, D2>,
79    <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
80{
81    change_layout_f(tensor, layout).map(|v| v.into_owned())
82}
83
84pub fn into_layout<'a, R, T, B, D, D2>(tensor: TensorAny<R, T, B, D>, layout: Layout<D2>) -> Tensor<T, B, D2>
85where
86    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
87    D: DimAPI,
88    D2: DimAPI,
89    T: Clone,
90    B: DeviceAPI<T>
91        + DeviceRawAPI<MaybeUninit<T>>
92        + DeviceCreationAnyAPI<T>
93        + OpAssignArbitaryAPI<T, D2, D>
94        + OpAssignAPI<T, D2>,
95    <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
96{
97    into_layout_f(tensor, layout).rstsr_unwrap()
98}
99
100pub fn change_layout<'a, R, T, B, D, D2>(tensor: TensorAny<R, T, B, D>, layout: Layout<D2>) -> TensorCow<'a, T, B, D2>
101where
102    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
103    D: DimAPI,
104    D2: DimAPI,
105    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D2, D>,
106{
107    change_layout_f(tensor, layout).rstsr_unwrap()
108}
109
110impl<'a, R, T, B, D> TensorAny<R, T, B, D>
111where
112    R: DataAPI<Data = B::Raw> + DataIntoCowAPI<'a>,
113    D: DimAPI,
114    T: Clone,
115    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
116{
117    /// Convert tensor to the other layout.
118    ///
119    /// # See also
120    ///
121    /// [`to_layout`]
122    pub fn to_layout<D2>(&self, layout: Layout<D2>) -> TensorCow<'_, T, B, D2>
123    where
124        D2: DimAPI,
125        B: OpAssignArbitaryAPI<T, D2, D>,
126    {
127        to_layout(self, layout)
128    }
129
130    pub fn to_layout_f<D2>(&self, layout: Layout<D2>) -> Result<TensorCow<'_, T, B, D2>>
131    where
132        D2: DimAPI,
133        B: OpAssignArbitaryAPI<T, D2, D>,
134    {
135        to_layout_f(self, layout)
136    }
137
138    pub fn into_layout_f<D2>(self, layout: Layout<D2>) -> Result<Tensor<T, B, D2>>
139    where
140        D2: DimAPI,
141        B: DeviceRawAPI<MaybeUninit<T>> + OpAssignArbitaryAPI<T, D2, D> + OpAssignAPI<T, D2>,
142        <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
143    {
144        into_layout_f(self, layout)
145    }
146
147    pub fn into_layout<D2>(self, layout: Layout<D2>) -> Tensor<T, B, D2>
148    where
149        D2: DimAPI,
150        B: DeviceRawAPI<MaybeUninit<T>> + OpAssignArbitaryAPI<T, D2, D> + OpAssignAPI<T, D2>,
151        <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
152    {
153        into_layout(self, layout)
154    }
155
156    pub fn change_layout_f<D2>(self, layout: Layout<D2>) -> Result<TensorCow<'a, T, B, D2>>
157    where
158        D2: DimAPI,
159        B: OpAssignArbitaryAPI<T, D2, D>,
160    {
161        change_layout_f(self, layout)
162    }
163
164    pub fn change_layout<D2>(self, layout: Layout<D2>) -> TensorCow<'a, T, B, D2>
165    where
166        D2: DimAPI,
167        B: OpAssignArbitaryAPI<T, D2, D>,
168    {
169        change_layout(self, layout)
170    }
171}
172
173/* #endregion */
rstsr_core/tensor/manuplication/to_layout.rs

rstsr_core/tensor/manuplication/
to_layout.rs
