rstsr_core/tensor/manuplication/
to_contig.rs

1use crate::prelude_dev::*;
2
3/* #region to_contig */
4
5/// Convert tensor to contiguous layout.
6///
7/// See also [`to_contig`].
8pub fn change_contig_f<'a, R, T, B, D>(
9    tensor: TensorAny<R, T, B, D>,
10    order: FlagOrder,
11) -> Result<TensorCow<'a, T, B, D>>
12where
13    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
14    D: DimAPI,
15    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
16{
17    let shape = tensor.shape();
18    let layout_new = match order {
19        RowMajor => shape.new_c_contig(None),
20        ColMajor => shape.new_f_contig(None),
21    };
22    change_layout_f(tensor, layout_new)
23}
24
25/// Convert tensor to contiguous layout.
26///
27/// This function takes a reference to a tensor and returns a [`TensorCow`] that is
28/// either a view (if the tensor is already contiguous with the requested order) or
29/// a newly allocated contiguous copy.
30///
31/// # Arguments
32///
33/// - `tensor`: A reference to the input tensor.
34/// - `order`: The memory layout order ([`RowMajor`] or [`ColMajor`]).
35///
36/// # Returns
37///
38/// A [`TensorCow`] containing either a view or an owned tensor with contiguous layout.
39///
40/// # Examples
41///
42/// ```rust
43/// # use rstsr::prelude::*;
44/// # let mut device = DeviceCpu::default();
45/// # device.set_default_order(RowMajor);
46/// // Create a non-contiguous tensor to contiguous
47/// let a = rt::arange((12, &device)).into_shape([3, 4]);
48/// let sliced = a.i((.., slice!(None, None, 2))); // Every other column
49/// println!("layout of sliced tensor: {:?}", sliced.layout());
50/// // 2-Dim (dyn), contiguous: Custom
51/// // shape: [3, 2], stride: [4, 2], offset: 0
52///
53/// // Convert to C-contiguous
54/// let contig = sliced.to_contig(RowMajor);
55/// println!("Contiguous layout: {:?}", contig.layout());
56/// // 2-Dim (dyn), contiguous: Cc
57/// // shape: [3, 2], stride: [2, 1], offset: 0
58/// ```
59///
60/// # See also
61///
62/// ## Similar functions in RSTSR
63///
64/// - [`reshape`]: Change the shape of a tensor without changing its data layout (returns
65///   copy-on-write tensor, view or necessarily clone). Reshape function inputs shape instead of
66///   layout.
67/// - [`to_prefer`]: Only converts if not already in preferred layout.
68///
69/// ## Variants of this function
70///
71/// - [`to_contig`] / [`to_contig_f`]: Non-consuming version that takes a reference and returns a
72///   view or owned tensor.
73/// - [`into_contig`] / [`into_contig_f`]: Consuming version that returns an owned tensor directly.
74/// - [`change_contig`] / [`change_contig_f`]: Consuming version that returns a view or owned
75///   tensor.
76/// - Associated methods on [`TensorAny`]:
77///
78///   - [`TensorAny::to_contig`] / [`TensorAny::to_contig_f`]
79///   - [`TensorAny::into_contig`] / [`TensorAny::into_contig_f`]
80///   - [`TensorAny::change_contig`] / [`TensorAny::change_contig_f`]
81pub fn to_contig<R, T, B, D>(tensor: &TensorAny<R, T, B, D>, order: FlagOrder) -> TensorCow<'_, T, B, D>
82where
83    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
84    D: DimAPI,
85    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
86{
87    to_contig_f(tensor, order).rstsr_unwrap()
88}
89
90/// Convert tensor to contiguous layout.
91///
92/// See also [`to_contig`].
93pub fn to_contig_f<R, T, B, D>(tensor: &TensorAny<R, T, B, D>, order: FlagOrder) -> Result<TensorCow<'_, T, B, D>>
94where
95    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
96    D: DimAPI,
97    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
98{
99    change_contig_f(tensor.view(), order)
100}
101
102/// Convert tensor to contiguous layout.
103///
104/// See also [`to_contig`].
105pub fn into_contig_f<'a, R, T, B, D>(tensor: TensorAny<R, T, B, D>, order: FlagOrder) -> Result<Tensor<T, B, D>>
106where
107    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
108    D: DimAPI,
109    T: Clone,
110    B: DeviceAPI<T>
111        + DeviceRawAPI<MaybeUninit<T>>
112        + DeviceCreationAnyAPI<T>
113        + OpAssignArbitaryAPI<T, D, D>
114        + OpAssignAPI<T, D>,
115    <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
116{
117    change_contig_f(tensor, order).map(|v| v.into_owned())
118}
119
120/// Convert tensor to contiguous layout.
121///
122/// See also [`to_contig`].
123pub fn change_contig<'a, R, T, B, D>(tensor: TensorAny<R, T, B, D>, order: FlagOrder) -> TensorCow<'a, T, B, D>
124where
125    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
126    D: DimAPI,
127    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
128{
129    change_contig_f(tensor, order).rstsr_unwrap()
130}
131
132/// Convert tensor to contiguous layout.
133///
134/// See also [`to_contig`].
135pub fn into_contig<'a, R, T, B, D>(tensor: TensorAny<R, T, B, D>, order: FlagOrder) -> Tensor<T, B, D>
136where
137    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
138    D: DimAPI,
139    T: Clone,
140    B: DeviceAPI<T>
141        + DeviceRawAPI<MaybeUninit<T>>
142        + DeviceCreationAnyAPI<T>
143        + OpAssignArbitaryAPI<T, D, D>
144        + OpAssignAPI<T, D>,
145    <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
146{
147    into_contig_f(tensor, order).rstsr_unwrap()
148}
149
150impl<'a, R, T, B, D> TensorAny<R, T, B, D>
151where
152    R: DataAPI<Data = B::Raw> + DataIntoCowAPI<'a>,
153    D: DimAPI,
154    T: Clone,
155    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
156{
157    /// Convert tensor to contiguous layout.
158    ///
159    /// See also [`to_contig`].
160    pub fn to_contig(&self, order: FlagOrder) -> TensorCow<'_, T, B, D>
161    where
162        B: OpAssignArbitaryAPI<T, D, D>,
163    {
164        to_contig(self, order)
165    }
166
167    /// Convert tensor to contiguous layout.
168    ///
169    /// See also [`to_contig`].
170    pub fn to_contig_f(&self, order: FlagOrder) -> Result<TensorCow<'_, T, B, D>>
171    where
172        B: OpAssignArbitaryAPI<T, D, D>,
173    {
174        to_contig_f(self, order)
175    }
176
177    /// Convert tensor to contiguous layout.
178    ///
179    /// See also [`to_contig`].
180    pub fn into_contig_f(self, order: FlagOrder) -> Result<Tensor<T, B, D>>
181    where
182        B: DeviceRawAPI<MaybeUninit<T>> + OpAssignArbitaryAPI<T, D, D> + OpAssignAPI<T, D>,
183        <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
184    {
185        into_contig_f(self, order)
186    }
187
188    /// Convert tensor to contiguous layout.
189    ///
190    /// # See also
191    ///
192    /// Refer to [`to_contig`] for more detailed documentation.
193    pub fn into_contig(self, order: FlagOrder) -> Tensor<T, B, D>
194    where
195        B: DeviceRawAPI<MaybeUninit<T>> + OpAssignArbitaryAPI<T, D, D> + OpAssignAPI<T, D>,
196        <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
197    {
198        into_contig(self, order)
199    }
200
201    /// Convert tensor to contiguous layout.
202    ///
203    /// See also [`to_contig`].
204    pub fn change_contig_f(self, order: FlagOrder) -> Result<TensorCow<'a, T, B, D>>
205    where
206        B: OpAssignArbitaryAPI<T, D, D>,
207    {
208        change_contig_f(self, order)
209    }
210
211    /// Convert tensor to contiguous layout.
212    ///
213    /// See also [`to_contig`].
214    pub fn change_contig(self, order: FlagOrder) -> TensorCow<'a, T, B, D>
215    where
216        B: OpAssignArbitaryAPI<T, D, D>,
217    {
218        change_contig(self, order)
219    }
220}
221
222/* #endregion */
223
224/* #region to_prefer */
225
226/// Convert tensor to preferred layout if not already contiguous.
227///
228/// See also [`to_prefer`].
229pub fn change_prefer_f<'a, R, T, B, D>(
230    tensor: TensorAny<R, T, B, D>,
231    order: FlagOrder,
232) -> Result<TensorCow<'a, T, B, D>>
233where
234    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
235    D: DimAPI,
236    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
237{
238    if (order == RowMajor && tensor.c_prefer()) || (order == ColMajor && tensor.f_prefer()) {
239        Ok(tensor.into_cow())
240    } else {
241        change_contig_f(tensor, order)
242    }
243}
244
245/// Convert tensor to preferred layout if not already contiguous.
246///
247/// This function checks if the tensor is already contiguous with the specified order.
248/// If it is, a view is returned without copying data. Otherwise, data is copied to
249/// a new contiguous layout.
250///
251/// # Arguments
252///
253/// - `tensor`: A reference to the input tensor.
254/// - `order`: The memory layout order ([`RowMajor`] or [`ColMajor`]).
255///
256/// # Returns
257///
258/// A [`TensorCow`] containing either a view (if already contiguous) or an owned tensor
259/// (if data was copied).
260///
261/// # Examples
262///
263/// ```rust
264/// # use rstsr::prelude::*;
265/// # let mut device = DeviceCpu::default();
266/// # device.set_default_order(RowMajor);
267/// // C-contiguous tensor stays as view
268/// let a = rt::tensor_from_nested!([[1, 2], [3, 4], [5, 6]], &device);
269/// let result = rt::to_prefer(&a, RowMajor);
270/// assert!(!result.is_owned());
271///
272/// // Transposed (non-contiguous) tensor gets copied
273/// let transposed = a.t();
274/// let result = rt::to_prefer(&transposed, RowMajor);
275/// assert!(result.is_owned());
276/// ```
277///
278/// # See also
279///
280/// ## Similar functions in RSTSR
281///
282/// - [`reshape`]: Change the shape of a tensor without changing its data layout (returns
283///   copy-on-write tensor, view or necessarily clone). Reshape function inputs shape instead of
284///   layout.
285/// - [`to_contig`]: Always converts to contiguous layout regardless of current layout.
286///
287/// ## Variants of this function
288///
289/// - [`to_prefer`] / [`to_prefer_f`]: Non-consuming version that takes a reference and returns a
290///   view or owned tensor.
291/// - [`into_prefer`] / [`into_prefer_f`]: Consuming version that returns an owned tensor directly.
292/// - [`change_prefer`] / [`change_prefer_f`]: Consuming version that returns a view or owned
293///   tensor.
294/// - Associated methods on [`TensorAny`]:
295///
296///   - [`TensorAny::to_prefer`] / [`TensorAny::to_prefer_f`]
297///   - [`TensorAny::into_prefer`] / [`TensorAny::into_prefer_f`]
298///   - [`TensorAny::change_prefer`] / [`TensorAny::change_prefer_f`]
299pub fn to_prefer<R, T, B, D>(tensor: &TensorAny<R, T, B, D>, order: FlagOrder) -> TensorCow<'_, T, B, D>
300where
301    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
302    D: DimAPI,
303    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
304{
305    to_prefer_f(tensor, order).rstsr_unwrap()
306}
307
308/// Convert tensor to preferred layout if not already contiguous.
309///
310/// See also [`to_prefer`].
311pub fn to_prefer_f<R, T, B, D>(tensor: &TensorAny<R, T, B, D>, order: FlagOrder) -> Result<TensorCow<'_, T, B, D>>
312where
313    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
314    D: DimAPI,
315    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
316{
317    change_prefer_f(tensor.view(), order)
318}
319
320/// Convert tensor to preferred layout if not already contiguous.
321///
322/// See also [`to_prefer`].
323pub fn into_prefer_f<'a, R, T, B, D>(tensor: TensorAny<R, T, B, D>, order: FlagOrder) -> Result<Tensor<T, B, D>>
324where
325    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
326    D: DimAPI,
327    T: Clone,
328    B: DeviceAPI<T>
329        + DeviceRawAPI<MaybeUninit<T>>
330        + DeviceCreationAnyAPI<T>
331        + OpAssignArbitaryAPI<T, D, D>
332        + OpAssignAPI<T, D>,
333    <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
334{
335    change_prefer_f(tensor, order).map(|v| v.into_owned())
336}
337
338/// Convert tensor to preferred layout if not already contiguous.
339///
340/// See also [`to_prefer`].
341pub fn change_prefer<'a, R, T, B, D>(tensor: TensorAny<R, T, B, D>, order: FlagOrder) -> TensorCow<'a, T, B, D>
342where
343    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
344    D: DimAPI,
345    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignArbitaryAPI<T, D, D>,
346{
347    change_prefer_f(tensor, order).rstsr_unwrap()
348}
349
350/// Convert tensor to preferred layout if not already contiguous.
351///
352/// See also [`to_prefer`].
353pub fn into_prefer<'a, R, T, B, D>(tensor: TensorAny<R, T, B, D>, order: FlagOrder) -> Tensor<T, B, D>
354where
355    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw> + DataIntoCowAPI<'a>,
356    D: DimAPI,
357    T: Clone,
358    B: DeviceAPI<T>
359        + DeviceRawAPI<MaybeUninit<T>>
360        + DeviceCreationAnyAPI<T>
361        + OpAssignArbitaryAPI<T, D, D>
362        + OpAssignAPI<T, D>,
363    <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
364{
365    into_prefer_f(tensor, order).rstsr_unwrap()
366}
367
368impl<'a, R, T, B, D> TensorAny<R, T, B, D>
369where
370    R: DataAPI<Data = B::Raw> + DataIntoCowAPI<'a>,
371    D: DimAPI,
372    T: Clone,
373    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
374{
375    /// Convert tensor to preferred layout if not already contiguous.
376    ///
377    /// See also [`to_prefer`].
378    pub fn to_prefer(&self, order: FlagOrder) -> TensorCow<'_, T, B, D>
379    where
380        B: OpAssignArbitaryAPI<T, D, D>,
381    {
382        to_prefer(self, order)
383    }
384
385    /// Convert tensor to preferred layout if not already contiguous.
386    ///
387    /// See also [`to_prefer`].
388    pub fn to_prefer_f(&self, order: FlagOrder) -> Result<TensorCow<'_, T, B, D>>
389    where
390        B: OpAssignArbitaryAPI<T, D, D>,
391    {
392        to_prefer_f(self, order)
393    }
394
395    /// Convert tensor to preferred layout if not already contiguous.
396    ///
397    /// See also [`to_prefer`].
398    pub fn into_prefer_f(self, order: FlagOrder) -> Result<Tensor<T, B, D>>
399    where
400        B: DeviceRawAPI<MaybeUninit<T>> + OpAssignArbitaryAPI<T, D, D> + OpAssignAPI<T, D>,
401        <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
402    {
403        into_prefer_f(self, order)
404    }
405
406    /// Convert tensor to preferred layout if not already contiguous.
407    ///
408    /// See also [`to_prefer`].
409    pub fn into_prefer(self, order: FlagOrder) -> Tensor<T, B, D>
410    where
411        B: DeviceRawAPI<MaybeUninit<T>> + OpAssignArbitaryAPI<T, D, D> + OpAssignAPI<T, D>,
412        <B as DeviceRawAPI<T>>::Raw: Clone + 'a,
413    {
414        into_prefer(self, order)
415    }
416
417    /// Convert tensor to preferred layout if not already contiguous.
418    ///
419    /// See also [`to_prefer`].
420    pub fn change_prefer_f(self, order: FlagOrder) -> Result<TensorCow<'a, T, B, D>>
421    where
422        B: OpAssignArbitaryAPI<T, D, D>,
423    {
424        change_prefer_f(self, order)
425    }
426
427    /// Convert tensor to preferred layout if not already contiguous.
428    ///
429    /// See also [`to_prefer`].
430    pub fn change_prefer(self, order: FlagOrder) -> TensorCow<'a, T, B, D>
431    where
432        B: OpAssignArbitaryAPI<T, D, D>,
433    {
434        change_prefer(self, order)
435    }
436}
437
438/* #endregion */
rstsr_core/tensor/manuplication/to_contig.rs

rstsr_core/tensor/manuplication/
to_contig.rs
