rstsr_core/tensor/
asarray.rs

1//! Implementation of function `asarray`.
2
3use crate::prelude_dev::*;
4use core::mem::ManuallyDrop;
5use num::complex::{Complex32, Complex64};
6
7pub trait AsArrayAPI<Inp> {
8    type Out;
9
10    fn asarray_f(self) -> Result<Self::Out>;
11
12    fn asarray(self) -> Self::Out
13    where
14        Self: Sized,
15    {
16        Self::asarray_f(self).rstsr_unwrap()
17    }
18}
19
20/// Convert the input to an array.
21///
22/// This function takes kinds of input and converts them to an array. Please
23/// refer to trait implementations of [`AsArrayAPI`].
24///
25/// # See also
26///
27/// [Python array API: `asarray`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.asarray.html)
28pub fn asarray<Args, Inp>(param: Args) -> Args::Out
29where
30    Args: AsArrayAPI<Inp>,
31{
32    return AsArrayAPI::asarray(param);
33}
34
35pub fn asarray_f<Args, Inp>(param: Args) -> Result<Args::Out>
36where
37    Args: AsArrayAPI<Inp>,
38{
39    return AsArrayAPI::asarray_f(param);
40}
41
42/* #region tensor input */
43
44impl<R, T, B, D> AsArrayAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder)
45where
46    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
47    T: Clone,
48    D: DimAPI,
49    B: DeviceAPI<T> + DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
50{
51    type Out = Tensor<T, B, D>;
52
53    fn asarray_f(self) -> Result<Self::Out> {
54        let (input, order) = self;
55        let device = input.device();
56        let layout_a = input.layout();
57        let layout_c = layout_for_array_copy(layout_a, order)?;
58        let mut storage_c = device.uninit_impl(layout_c.size())?;
59        device.assign_uninit(storage_c.raw_mut(), &layout_c, input.raw(), layout_a)?;
60        let storage_c = unsafe { B::assume_init_impl(storage_c) }?;
61        let tensor = unsafe { Tensor::new_unchecked(storage_c, layout_c) };
62        return Ok(tensor);
63    }
64}
65
66impl<R, T, B, D> AsArrayAPI<()> for &TensorAny<R, T, B, D>
67where
68    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
69    T: Clone,
70    D: DimAPI,
71    B: DeviceAPI<T> + DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
72{
73    type Out = Tensor<T, B, D>;
74
75    fn asarray_f(self) -> Result<Self::Out> {
76        asarray_f((self, TensorIterOrder::default()))
77    }
78}
79
80impl<T, B, D> AsArrayAPI<()> for (Tensor<T, B, D>, TensorIterOrder)
81where
82    T: Clone,
83    D: DimAPI,
84    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
85{
86    type Out = Tensor<T, B, D>;
87
88    fn asarray_f(self) -> Result<Self::Out> {
89        let (input, order) = self;
90        let storage_a = input.storage();
91        let layout_a = input.layout();
92        let device = storage_a.device();
93        let layout_c = layout_for_array_copy(layout_a, order)?;
94        if layout_c == *layout_a {
95            return Ok(input);
96        } else {
97            let mut storage_c = device.uninit_impl(layout_c.size())?;
98            device.assign_uninit(storage_c.raw_mut(), &layout_c, storage_a.raw(), layout_a)?;
99            let storage_c = unsafe { B::assume_init_impl(storage_c) }?;
100            let tensor = unsafe { Tensor::new_unchecked(storage_c, layout_c) };
101            return Ok(tensor);
102        }
103    }
104}
105
106impl<T, B, D> AsArrayAPI<()> for Tensor<T, B, D>
107where
108    T: Clone,
109    D: DimAPI,
110    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
111{
112    type Out = Tensor<T, B, D>;
113
114    fn asarray_f(self) -> Result<Self::Out> {
115        asarray_f((self, TensorIterOrder::default()))
116    }
117}
118
119/* #endregion */
120
121/* #region vec-like input */
122
123impl<T, B> AsArrayAPI<()> for (Vec<T>, &B)
124where
125    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
126{
127    type Out = Tensor<T, B, IxD>;
128
129    fn asarray_f(self) -> Result<Self::Out> {
130        let (input, device) = self;
131        let layout = vec![input.len()].c();
132        let storage = device.outof_cpu_vec(input)?;
133        let tensor = unsafe { Tensor::new_unchecked(storage, layout) };
134        return Ok(tensor);
135    }
136}
137
138impl<T, B, D> AsArrayAPI<D> for (Vec<T>, Layout<D>, &B)
139where
140    D: DimAPI,
141    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
142{
143    type Out = Tensor<T, B, IxD>;
144
145    fn asarray_f(self) -> Result<Self::Out> {
146        let (input, layout, device) = self;
147        rstsr_assert_eq!(
148            layout.bounds_index()?,
149            (0, layout.size()),
150            InvalidLayout,
151            "This constructor assumes compact memory layout."
152        )?;
153        rstsr_assert_eq!(
154            layout.size(),
155            input.len(),
156            InvalidLayout,
157            "This constructor assumes that the layout size is equal to the input size."
158        )?;
159        let storage = device.outof_cpu_vec(input)?;
160        let tensor = unsafe { Tensor::new_unchecked(storage, layout.into_dim()?) };
161        return Ok(tensor);
162    }
163}
164
165impl<T, B, D> AsArrayAPI<D> for (Vec<T>, D, &B)
166where
167    D: DimAPI,
168    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
169{
170    type Out = Tensor<T, B, IxD>;
171
172    fn asarray_f(self) -> Result<Self::Out> {
173        let (input, shape, device) = self;
174        let default_order = device.default_order();
175        let layout = match default_order {
176            RowMajor => shape.c(),
177            ColMajor => shape.f(),
178        };
179        asarray_f((input, layout, device))
180    }
181}
182
183impl<T> AsArrayAPI<()> for Vec<T>
184where
185    T: Clone,
186{
187    type Out = Tensor<T, DeviceCpu, IxD>;
188
189    fn asarray_f(self) -> Result<Self::Out> {
190        asarray_f((self, &DeviceCpu::default()))
191    }
192}
193
194#[duplicate_item(L; [D]; [Layout<D>])]
195impl<T, D> AsArrayAPI<D> for (Vec<T>, L)
196where
197    T: Clone,
198    D: DimAPI,
199{
200    type Out = Tensor<T, DeviceCpu, IxD>;
201
202    fn asarray_f(self) -> Result<Self::Out> {
203        let (input, layout) = self;
204        asarray_f((input, layout, &DeviceCpu::default()))
205    }
206}
207
208impl<T> From<Vec<T>> for Tensor<T, DeviceCpu, IxD>
209where
210    T: Clone,
211{
212    fn from(input: Vec<T>) -> Self {
213        asarray_f(input).rstsr_unwrap()
214    }
215}
216
217/* #endregion */
218
219/* #region slice-like input */
220
221impl<'a, T, B, D> AsArrayAPI<D> for (&'a [T], Layout<D>, &B)
222where
223    T: Clone,
224    B: DeviceAPI<T, Raw = Vec<T>>,
225    D: DimAPI,
226{
227    type Out = TensorView<'a, T, B, IxD>;
228
229    fn asarray_f(self) -> Result<Self::Out> {
230        let (input, layout, device) = self;
231        rstsr_assert_eq!(
232            layout.bounds_index()?,
233            (0, layout.size()),
234            InvalidLayout,
235            "This constructor assumes compact memory layout."
236        )?;
237        rstsr_assert_eq!(
238            layout.size(),
239            input.len(),
240            InvalidLayout,
241            "This constructor assumes that the layout size is equal to the input size."
242        )?;
243        let ptr = input.as_ptr();
244        let len = input.len();
245        let raw = unsafe {
246            let ptr = ptr as *mut T;
247            Vec::from_raw_parts(ptr, len, len)
248        };
249        let device = device.clone();
250        let data = DataRef::from_manually_drop(ManuallyDrop::new(raw));
251        let storage = Storage::new(data, device);
252        let tensor = unsafe { TensorView::new_unchecked(storage, layout.into_dim()?) };
253        return Ok(tensor);
254    }
255}
256
257impl<'a, T, B, D> AsArrayAPI<D> for (&'a [T], D, &B)
258where
259    T: Clone,
260    B: DeviceAPI<T, Raw = Vec<T>>,
261    D: DimAPI,
262{
263    type Out = TensorView<'a, T, B, IxD>;
264
265    fn asarray_f(self) -> Result<Self::Out> {
266        let (input, shape, device) = self;
267        let default_order = device.default_order();
268        let layout = match default_order {
269            RowMajor => shape.c(),
270            ColMajor => shape.f(),
271        };
272        asarray_f((input, layout, device))
273    }
274}
275
276impl<'a, T, B> AsArrayAPI<()> for (&'a [T], &B)
277where
278    T: Clone,
279    B: DeviceAPI<T, Raw = Vec<T>>,
280{
281    type Out = TensorView<'a, T, B, IxD>;
282
283    fn asarray_f(self) -> Result<Self::Out> {
284        let (input, device) = self;
285        let layout = vec![input.len()].c();
286        let device = device.clone();
287
288        let ptr = input.as_ptr();
289        let len = input.len();
290        let raw = unsafe {
291            let ptr = ptr as *mut T;
292            Vec::from_raw_parts(ptr, len, len)
293        };
294        let data = DataRef::from_manually_drop(ManuallyDrop::new(raw));
295        let storage = Storage::new(data, device);
296        let tensor = unsafe { TensorView::new_unchecked(storage, layout) };
297        return Ok(tensor);
298    }
299}
300
301#[duplicate_item(L; [D]; [Layout<D>])]
302impl<'a, T, D> AsArrayAPI<D> for (&'a [T], L)
303where
304    T: Clone,
305    D: DimAPI,
306{
307    type Out = TensorView<'a, T, DeviceCpu, IxD>;
308
309    fn asarray_f(self) -> Result<Self::Out> {
310        let (input, layout) = self;
311        asarray_f((input, layout, &DeviceCpu::default()))
312    }
313}
314
315impl<'a, T> AsArrayAPI<()> for &'a [T]
316where
317    T: Clone,
318{
319    type Out = TensorView<'a, T, DeviceCpu, IxD>;
320
321    fn asarray_f(self) -> Result<Self::Out> {
322        asarray_f((self, &DeviceCpu::default()))
323    }
324}
325
326#[duplicate_item(L; [D]; [Layout<D>])]
327impl<'a, T, B, D> AsArrayAPI<D> for (&'a Vec<T>, L, &B)
328where
329    T: Clone,
330    B: DeviceAPI<T, Raw = Vec<T>> + 'a,
331    D: DimAPI,
332{
333    type Out = TensorView<'a, T, B, IxD>;
334
335    fn asarray_f(self) -> Result<Self::Out> {
336        let (input, layout, device) = self;
337        asarray_f((input.as_slice(), layout, device))
338    }
339}
340
341impl<'a, T, B> AsArrayAPI<()> for (&'a Vec<T>, &B)
342where
343    T: Clone,
344    B: DeviceAPI<T, Raw = Vec<T>>,
345{
346    type Out = TensorView<'a, T, B, IxD>;
347
348    fn asarray_f(self) -> Result<Self::Out> {
349        let (input, device) = self;
350        asarray_f((input.as_slice(), device))
351    }
352}
353
354#[duplicate_item(L; [D]; [Layout<D>])]
355impl<'a, T, D> AsArrayAPI<D> for (&'a Vec<T>, L)
356where
357    T: Clone,
358    D: DimAPI,
359{
360    type Out = TensorView<'a, T, DeviceCpu, IxD>;
361
362    fn asarray_f(self) -> Result<Self::Out> {
363        let (input, layout) = self;
364        asarray_f((input.as_slice(), layout, &DeviceCpu::default()))
365    }
366}
367
368impl<'a, T> AsArrayAPI<()> for &'a Vec<T>
369where
370    T: Clone,
371{
372    type Out = TensorView<'a, T, DeviceCpu, IxD>;
373
374    fn asarray_f(self) -> Result<Self::Out> {
375        asarray_f((self.as_slice(), &DeviceCpu::default()))
376    }
377}
378
379impl<'a, T> From<&'a [T]> for TensorView<'a, T, DeviceCpu, IxD>
380where
381    T: Clone,
382{
383    fn from(input: &'a [T]) -> Self {
384        asarray(input)
385    }
386}
387
388impl<'a, T> From<&'a Vec<T>> for TensorView<'a, T, DeviceCpu, IxD>
389where
390    T: Clone,
391{
392    fn from(input: &'a Vec<T>) -> Self {
393        asarray(input)
394    }
395}
396
397/* #endregion */
398
399/* #region slice-like mutable input */
400
401impl<'a, T, B, D> AsArrayAPI<D> for (&'a mut [T], Layout<D>, &B)
402where
403    T: Clone,
404    B: DeviceAPI<T, Raw = Vec<T>>,
405    D: DimAPI,
406{
407    type Out = TensorMut<'a, T, B, IxD>;
408
409    fn asarray_f(self) -> Result<Self::Out> {
410        let (input, layout, device) = self;
411        rstsr_assert_eq!(
412            layout.bounds_index()?,
413            (0, layout.size()),
414            InvalidLayout,
415            "This constructor assumes compact memory layout."
416        )?;
417        rstsr_assert_eq!(
418            layout.size(),
419            input.len(),
420            InvalidLayout,
421            "This constructor assumes that the layout size is equal to the input size."
422        )?;
423        let ptr = input.as_ptr();
424        let len = input.len();
425        let raw = unsafe {
426            let ptr = ptr as *mut T;
427            Vec::from_raw_parts(ptr, len, len)
428        };
429        let device = device.clone();
430        let data = DataMut::from_manually_drop(ManuallyDrop::new(raw));
431        let storage = Storage::new(data, device);
432        let tensor = unsafe { TensorMut::new_unchecked(storage, layout.into_dim()?) };
433        return Ok(tensor);
434    }
435}
436
437impl<'a, T, B, D> AsArrayAPI<D> for (&'a mut [T], D, &B)
438where
439    T: Clone,
440    B: DeviceAPI<T, Raw = Vec<T>>,
441    D: DimAPI,
442{
443    type Out = TensorMut<'a, T, B, IxD>;
444
445    fn asarray_f(self) -> Result<Self::Out> {
446        let (input, shape, device) = self;
447        let default_order = device.default_order();
448        let layout = match default_order {
449            RowMajor => shape.c(),
450            ColMajor => shape.f(),
451        };
452        asarray_f((input, layout, device))
453    }
454}
455
456impl<'a, T, B> AsArrayAPI<()> for (&'a mut [T], &B)
457where
458    T: Clone,
459    B: DeviceAPI<T, Raw = Vec<T>>,
460{
461    type Out = TensorMut<'a, T, B, IxD>;
462
463    fn asarray_f(self) -> Result<Self::Out> {
464        let (input, device) = self;
465        let layout = [input.len()].c();
466        let device = device.clone();
467
468        let ptr = input.as_ptr();
469        let len = input.len();
470        let raw = unsafe {
471            let ptr = ptr as *mut T;
472            Vec::from_raw_parts(ptr, len, len)
473        };
474        let data = DataMut::from_manually_drop(ManuallyDrop::new(raw));
475        let storage = Storage::new(data, device);
476        let tensor = unsafe { TensorMut::new_unchecked(storage, layout.into_dim()?) };
477        return Ok(tensor);
478    }
479}
480
481#[duplicate_item(L; [D]; [Layout<D>])]
482impl<'a, T, D> AsArrayAPI<D> for (&'a mut [T], L)
483where
484    T: Clone,
485    D: DimAPI,
486{
487    type Out = TensorMut<'a, T, DeviceCpu, IxD>;
488
489    fn asarray_f(self) -> Result<Self::Out> {
490        let (input, layout) = self;
491        asarray_f((input, layout, &DeviceCpu::default()))
492    }
493}
494
495impl<'a, T> AsArrayAPI<()> for &'a mut [T]
496where
497    T: Clone,
498{
499    type Out = TensorMut<'a, T, DeviceCpu, IxD>;
500
501    fn asarray_f(self) -> Result<Self::Out> {
502        asarray_f((self, &DeviceCpu::default()))
503    }
504}
505
506#[duplicate_item(L; [D]; [Layout<D>])]
507impl<'a, T, B, D> AsArrayAPI<D> for (&'a mut Vec<T>, L, &B)
508where
509    T: Clone,
510    B: DeviceAPI<T, Raw = Vec<T>>,
511    D: DimAPI,
512{
513    type Out = TensorMut<'a, T, B, IxD>;
514
515    fn asarray_f(self) -> Result<Self::Out> {
516        let (input, layout, device) = self;
517        asarray_f((input.as_mut_slice(), layout, device))
518    }
519}
520
521impl<'a, T, B> AsArrayAPI<()> for (&'a mut Vec<T>, &B)
522where
523    T: Clone,
524    B: DeviceAPI<T, Raw = Vec<T>>,
525{
526    type Out = TensorMut<'a, T, B, IxD>;
527
528    fn asarray_f(self) -> Result<Self::Out> {
529        let (input, device) = self;
530        asarray_f((input.as_mut_slice(), device))
531    }
532}
533
534#[duplicate_item(L; [D]; [Layout<D>])]
535impl<'a, T, D> AsArrayAPI<D> for (&'a mut Vec<T>, L)
536where
537    T: Clone,
538    D: DimAPI,
539{
540    type Out = TensorMut<'a, T, DeviceCpu, IxD>;
541
542    fn asarray_f(self) -> Result<Self::Out> {
543        let (input, layout) = self;
544        asarray_f((input.as_mut_slice(), layout, &DeviceCpu::default()))
545    }
546}
547
548impl<'a, T> AsArrayAPI<()> for &'a mut Vec<T>
549where
550    T: Clone,
551{
552    type Out = TensorMut<'a, T, DeviceCpu, IxD>;
553
554    fn asarray_f(self) -> Result<Self::Out> {
555        asarray_f((self.as_mut_slice(), &DeviceCpu::default()))
556    }
557}
558
559impl<'a, T> From<&'a mut [T]> for TensorMut<'a, T, DeviceCpu, IxD>
560where
561    T: Clone,
562{
563    fn from(input: &'a mut [T]) -> Self {
564        asarray(input)
565    }
566}
567
568impl<'a, T> From<&'a mut Vec<T>> for TensorMut<'a, T, DeviceCpu, IxD>
569where
570    T: Clone,
571{
572    fn from(input: &'a mut Vec<T>) -> Self {
573        asarray(input)
574    }
575}
576
577/* #endregion */
578
579/* #region scalar input */
580
581macro_rules! impl_asarray_scalar {
582    ($($t:ty),*) => {
583        $(
584            impl<B> AsArrayAPI<()> for ($t, &B)
585            where
586                B: DeviceAPI<$t> + DeviceCreationAnyAPI<$t>,
587            {
588                type Out = Tensor<$t, B, IxD>;
589
590                fn asarray_f(self) -> Result<Self::Out> {
591                    let (input, device) = self;
592                    let layout = Layout::new(vec![], vec![], 0)?;
593                    let storage = device.outof_cpu_vec(vec![input])?;
594                    let tensor = unsafe { Tensor::new_unchecked(storage, layout) };
595                    return Ok(tensor);
596                }
597            }
598
599            impl AsArrayAPI<()> for $t {
600                type Out = Tensor<$t, DeviceCpu, IxD>;
601
602                fn asarray_f(self) -> Result<Self::Out> {
603                    asarray_f((self, &DeviceCpu::default()))
604                }
605            }
606        )*
607    };
608}
609
610impl_asarray_scalar!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize, f32, f64, Complex32, Complex64);
611
612/* #endregion */
613
614#[cfg(test)]
615mod tests {
616    use super::*;
617
618    #[test]
619    fn test_asarray() {
620        let input = vec![1, 2, 3];
621        let tensor = asarray_f(input).unwrap();
622        println!("{tensor:?}");
623        let input = [1, 2, 3];
624        let tensor = asarray_f(input.as_ref()).unwrap();
625        println!("{tensor:?}");
626
627        let input = vec![1, 2, 3];
628        let tensor = asarray_f(&input).unwrap();
629        println!("{:?}", tensor.raw().as_ptr());
630        println!("{tensor:?}");
631
632        let tensor = asarray_f((&tensor, TensorIterOrder::K)).unwrap();
633        println!("{tensor:?}");
634
635        let tensor = asarray_f((tensor, TensorIterOrder::K)).unwrap();
636        println!("{tensor:?}");
637    }
638
639    #[test]
640    fn test_asarray_scalar() {
641        let tensor = asarray_f(1).unwrap();
642        println!("{tensor:?}");
643        let tensor = asarray_f((Complex64::new(0., 1.), &DeviceCpuSerial::default())).unwrap();
644        println!("{tensor:?}");
645    }
646}
rstsr_core/tensor/asarray.rs

rstsr_core/tensor/
asarray.rs
