rstsr_core/tensor/
creation.rs

1//! Creation methods for `Tensor` struct.
2//!
3//! This module relates to the [Python array API standard v2024.12](https://data-apis.org/array-api/2024.12/API_specification/creation_functions.html).
4//!
5//! Todo list:
6//! - [x] `arange`: [`arange`]
7//! - [x] `asarray`: [`asarray`] (defined elsewhere)
8//! - [x] `empty`: [`empty`]
9//! - [x] `empty_like`: [`empty_like`]
10//! - [x] `eye`: [`eye`]
11//! - [ ] ~`from_dlpack`~
12//! - [x] `full`: [`full`]
13//! - [x] `full_like`: [`full_like`]
14//! - [x] `linspace`: [`linspace`]
15//! - [x] `meshgrid`
16//! - [x] `ones`: [`ones`]
17//! - [x] `ones_like`: [`ones_like`]
18//! - [x] `tril`
19//! - [x] `triu`
20//! - [x] `zeros`: [`zeros`]
21//! - [x] `zeros_like`: [`zeros_like`]
22
23use crate::prelude_dev::*;
24use num::complex::ComplexFloat;
25use num::Num;
26
27/* #region arange */
28
29pub trait ArangeAPI<Inp> {
30    type Out;
31
32    fn arange_f(self) -> Result<Self::Out>;
33
34    fn arange(self) -> Self::Out
35    where
36        Self: Sized,
37    {
38        Self::arange_f(self).rstsr_unwrap()
39    }
40}
41
42/// Evenly spaced values within the half-open interval `[start, stop)` as
43/// one-dimensional array.
44///
45/// # See also
46///
47/// - [Python array API standard: `arange`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.arange.html)
48pub fn arange<Args, Inp>(param: Args) -> Args::Out
49where
50    Args: ArangeAPI<Inp>,
51{
52    return ArangeAPI::arange(param);
53}
54
55pub fn arange_f<Args, Inp>(param: Args) -> Result<Args::Out>
56where
57    Args: ArangeAPI<Inp>,
58{
59    return ArangeAPI::arange_f(param);
60}
61
62impl<T, B> ArangeAPI<(T, B)> for (T, T, T, &B)
63where
64    T: Num + PartialOrd,
65    B: DeviceAPI<T> + DeviceCreationPartialOrdNumAPI<T>,
66{
67    type Out = Tensor<T, B, IxD>;
68
69    fn arange_f(self) -> Result<Self::Out> {
70        // full implementation
71        let (start, stop, step, device) = self;
72        let data = device.arange_impl(start, stop, step)?;
73        let layout = vec![data.len()].into();
74        unsafe { Ok(Tensor::new_unchecked(data, layout)) }
75    }
76}
77
78impl<T, B> ArangeAPI<(T, B)> for (T, T, &B)
79where
80    T: Num + PartialOrd,
81    B: DeviceAPI<T> + DeviceCreationPartialOrdNumAPI<T>,
82{
83    type Out = Tensor<T, B, IxD>;
84
85    fn arange_f(self) -> Result<Self::Out> {
86        // (start, stop, device) -> (start, stop, 1, device)
87        let (start, stop, device) = self;
88        let step = T::one();
89        arange_f((start, stop, step, device))
90    }
91}
92
93impl<T, B> ArangeAPI<(T, B)> for (T, &B)
94where
95    T: Num + PartialOrd,
96    B: DeviceAPI<T> + DeviceCreationPartialOrdNumAPI<T>,
97{
98    type Out = Tensor<T, B, IxD>;
99
100    fn arange_f(self) -> Result<Self::Out> {
101        // (stop, device) -> (0, stop, 1, device)
102        let (stop, device) = self;
103        let start = T::zero();
104        let step = T::one();
105        arange_f((start, stop, step, device))
106    }
107}
108
109impl<T> ArangeAPI<T> for (T, T, T)
110where
111    T: Num + PartialOrd + Clone + Send + Sync,
112{
113    type Out = Tensor<T, DeviceCpu, IxD>;
114
115    fn arange_f(self) -> Result<Self::Out> {
116        // full implementation
117        let (start, stop, step) = self;
118        arange_f((start, stop, step, &DeviceCpu::default()))
119    }
120}
121
122impl<T> ArangeAPI<T> for (T, T)
123where
124    T: Num + PartialOrd + Clone,
125{
126    type Out = Tensor<T, DeviceCpu, IxD>;
127
128    fn arange_f(self) -> Result<Self::Out> {
129        // (start, stop) -> (start, stop, 1)
130        let (start, stop) = self;
131        arange_f((start, stop, &DeviceCpu::default()))
132    }
133}
134
135impl<T> ArangeAPI<T> for T
136where
137    T: Num + PartialOrd + Clone,
138{
139    type Out = Tensor<T, DeviceCpu, IxD>;
140
141    fn arange_f(self) -> Result<Self::Out> {
142        // (stop) -> (0, stop, 1)
143        arange_f((T::zero(), self, &DeviceCpu::default()))
144    }
145}
146
147/* #endregion */
148
149/* #region empty */
150
151pub trait EmptyAPI<Inp> {
152    type Out;
153
154    /// # Safety
155    ///
156    /// This function is unsafe because it creates a tensor with uninitialized.
157    unsafe fn empty_f(self) -> Result<Self::Out>;
158
159    /// # Safety
160    ///
161    /// This function is unsafe because it creates a tensor with uninitialized.
162    unsafe fn empty(self) -> Self::Out
163    where
164        Self: Sized,
165    {
166        Self::empty_f(self).rstsr_unwrap()
167    }
168}
169
170/// Uninitialized tensor having a specified shape.
171///
172/// # Safety
173///
174/// This function is unsafe because it creates a tensor with uninitialized.
175///
176/// # See also
177///
178/// - [Python array API standard: `empty`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.empty.html)
179pub unsafe fn empty<Args, Inp>(param: Args) -> Args::Out
180where
181    Args: EmptyAPI<Inp>,
182{
183    return EmptyAPI::empty(param);
184}
185
186/// # Safety
187///
188/// This function is unsafe because it creates a tensor with uninitialized.
189pub unsafe fn empty_f<Args, Inp>(param: Args) -> Result<Args::Out>
190where
191    Args: EmptyAPI<Inp>,
192{
193    return EmptyAPI::empty_f(param);
194}
195
196impl<T, D, B> EmptyAPI<(T, D)> for (Layout<D>, &B)
197where
198    D: DimAPI,
199    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
200{
201    type Out = Tensor<T, B, IxD>;
202
203    unsafe fn empty_f(self) -> Result<Self::Out> {
204        let (layout, device) = self;
205        let (_, idx_max) = layout.bounds_index()?;
206        let storage = B::empty_impl(device, idx_max)?;
207        unsafe { Ok(Tensor::new_unchecked(storage, layout.into_dim()?)) }
208    }
209}
210
211impl<T, D, B> EmptyAPI<(T, D)> for (D, FlagOrder, &B)
212where
213    D: DimAPI,
214    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
215{
216    type Out = Tensor<T, B, IxD>;
217
218    unsafe fn empty_f(self) -> Result<Self::Out> {
219        let (shape, order, device) = self;
220        let layout = shape.new_contig(None, order);
221        empty_f((layout, device))
222    }
223}
224
225impl<T, D, B> EmptyAPI<(T, D)> for (D, &B)
226where
227    D: DimAPI,
228    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
229{
230    type Out = Tensor<T, B, IxD>;
231
232    unsafe fn empty_f(self) -> Result<Self::Out> {
233        let (shape, device) = self;
234        let default_order = device.default_order();
235        let layout = shape.new_contig(None, default_order);
236        empty_f((layout, device))
237    }
238}
239
240impl<T, D> EmptyAPI<(T, D)> for (D, FlagOrder)
241where
242    D: DimAPI,
243{
244    type Out = Tensor<T, DeviceCpu, IxD>;
245
246    unsafe fn empty_f(self) -> Result<Self::Out> {
247        let (shape, order) = self;
248        empty_f((shape, order, &DeviceCpu::default()))
249    }
250}
251
252#[duplicate_item(L; [D]; [Layout<D>])]
253impl<T, D> EmptyAPI<(T, D)> for L
254where
255    D: DimAPI,
256{
257    type Out = Tensor<T, DeviceCpu, IxD>;
258
259    unsafe fn empty_f(self) -> Result<Self::Out> {
260        empty_f((self, &DeviceCpu::default()))
261    }
262}
263
264/* #endregion */
265
266/* #region empty_like */
267
268pub trait EmptyLikeAPI<Inp> {
269    type Out;
270
271    /// # Safety
272    ///
273    /// This function is unsafe because it creates a tensor with uninitialized.
274    unsafe fn empty_like_f(self) -> Result<Self::Out>;
275
276    /// # Safety
277    ///
278    /// This function is unsafe because it creates a tensor with uninitialized.
279    unsafe fn empty_like(self) -> Self::Out
280    where
281        Self: Sized,
282    {
283        Self::empty_like_f(self).rstsr_unwrap()
284    }
285}
286
287/// Uninitialized tensor with the same shape as an input tensor.
288///
289/// # Safety
290///
291/// This function is unsafe because it creates a tensor with uninitialized.
292///
293/// # See also
294///
295/// - [Python array API standard: `empty_like`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.empty_like.html)
296pub unsafe fn empty_like<Args, Inp>(param: Args) -> Args::Out
297where
298    Args: EmptyLikeAPI<Inp>,
299{
300    return EmptyLikeAPI::empty_like(param);
301}
302
303/// # Safety
304///
305/// This function is unsafe because it creates a tensor with uninitialized.
306pub unsafe fn empty_like_f<Args, Inp>(param: Args) -> Result<Args::Out>
307where
308    Args: EmptyLikeAPI<Inp>,
309{
310    return EmptyLikeAPI::empty_like_f(param);
311}
312
313impl<R, T, B, D> EmptyLikeAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder, &B)
314where
315    R: DataAPI<Data = B::Raw>,
316    D: DimAPI,
317    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
318{
319    type Out = Tensor<T, B, D>;
320
321    unsafe fn empty_like_f(self) -> Result<Self::Out> {
322        let (tensor, order, device) = self;
323        let layout = layout_for_array_copy(tensor.layout(), order)?;
324        let idx_max = layout.size();
325        let storage = device.empty_impl(idx_max)?;
326        unsafe { Ok(Tensor::new_unchecked(storage, layout)) }
327    }
328}
329
330impl<R, T, B, D> EmptyLikeAPI<()> for (&TensorAny<R, T, B, D>, &B)
331where
332    R: DataAPI<Data = B::Raw>,
333    D: DimAPI,
334    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
335{
336    type Out = Tensor<T, B, D>;
337
338    unsafe fn empty_like_f(self) -> Result<Self::Out> {
339        let (tensor, device) = self;
340        empty_like_f((tensor, TensorIterOrder::default(), device))
341    }
342}
343
344impl<R, T, B, D> EmptyLikeAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder)
345where
346    R: DataAPI<Data = B::Raw>,
347    D: DimAPI,
348    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
349{
350    type Out = Tensor<T, B, D>;
351
352    unsafe fn empty_like_f(self) -> Result<Self::Out> {
353        let (tensor, order) = self;
354        let device = tensor.device();
355        empty_like_f((tensor, order, device))
356    }
357}
358
359impl<R, T, B, D> EmptyLikeAPI<()> for &TensorAny<R, T, B, D>
360where
361    R: DataAPI<Data = B::Raw>,
362    D: DimAPI,
363    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
364{
365    type Out = Tensor<T, B, D>;
366
367    unsafe fn empty_like_f(self) -> Result<Self::Out> {
368        let device = self.device();
369        empty_like_f((self, TensorIterOrder::default(), device))
370    }
371}
372
373/* #endregion */
374
375/* #region eye */
376
377pub trait EyeAPI<Inp> {
378    type Out;
379
380    fn eye_f(self) -> Result<Self::Out>;
381
382    fn eye(self) -> Self::Out
383    where
384        Self: Sized,
385    {
386        Self::eye_f(self).rstsr_unwrap()
387    }
388}
389
390/// Returns a two-dimensional array with ones on the kth diagonal and zeros
391/// elsewhere.
392///
393/// # See also
394///
395/// - [Python array API standard: `eye`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.eye.html)
396pub fn eye<Args, Inp>(param: Args) -> Args::Out
397where
398    Args: EyeAPI<Inp>,
399{
400    return EyeAPI::eye(param);
401}
402
403pub fn eye_f<Args, Inp>(param: Args) -> Result<Args::Out>
404where
405    Args: EyeAPI<Inp>,
406{
407    return EyeAPI::eye_f(param);
408}
409
410impl<T, B> EyeAPI<(T, B)> for (usize, usize, isize, FlagOrder, &B)
411where
412    T: Num,
413    B: DeviceAPI<T> + DeviceCreationNumAPI<T> + OpAssignAPI<T, Ix1>,
414{
415    type Out = Tensor<T, B, IxD>;
416
417    fn eye_f(self) -> Result<Self::Out> {
418        let (n_rows, n_cols, k, order, device) = self;
419        let layout = match order {
420            RowMajor => [n_rows, n_cols].c(),
421            ColMajor => [n_cols, n_rows].f(),
422        };
423        let mut storage = device.zeros_impl(layout.size())?;
424        let layout_diag = layout.diagonal(Some(k), Some(0), Some(1))?;
425        device.fill(storage.raw_mut(), &layout_diag, T::one())?;
426        unsafe { Ok(Tensor::new_unchecked(storage, layout.into_dim()?)) }
427    }
428}
429
430impl<T, B> EyeAPI<(T, B)> for (usize, usize, isize, &B)
431where
432    T: Num,
433    B: DeviceAPI<T> + DeviceCreationNumAPI<T> + OpAssignAPI<T, Ix1>,
434{
435    type Out = Tensor<T, B, IxD>;
436
437    fn eye_f(self) -> Result<Self::Out> {
438        // (n_rows, n_cols, k, device) -> (n_rows, n_cols, k, C, device)
439        let (n_rows, n_cols, k, device) = self;
440        let default_order = device.default_order();
441        eye_f((n_rows, n_cols, k, default_order, device))
442    }
443}
444
445impl<T, B> EyeAPI<(T, B)> for (usize, &B)
446where
447    T: Num,
448    B: DeviceAPI<T> + DeviceCreationNumAPI<T> + OpAssignAPI<T, Ix1>,
449{
450    type Out = Tensor<T, B, IxD>;
451
452    fn eye_f(self) -> Result<Self::Out> {
453        // (n_rows, n_cols, k, device) -> (n_rows, n_cols, k, C, device)
454        let (n_rows, device) = self;
455        let default_order = device.default_order();
456        eye_f((n_rows, n_rows, 0, default_order, device))
457    }
458}
459
460impl<T> EyeAPI<T> for (usize, usize, isize, FlagOrder)
461where
462    T: Num + Clone + Send + Sync,
463{
464    type Out = Tensor<T, DeviceCpu, IxD>;
465
466    fn eye_f(self) -> Result<Self::Out> {
467        let (n_rows, n_cols, k, order) = self;
468        eye_f((n_rows, n_cols, k, order, &DeviceCpu::default()))
469    }
470}
471
472impl<T> EyeAPI<T> for (usize, usize, isize)
473where
474    T: Num + Clone + Send + Sync,
475{
476    type Out = Tensor<T, DeviceCpu, IxD>;
477
478    fn eye_f(self) -> Result<Self::Out> {
479        // (n_rows, n_cols, k) -> (n_rows, n_cols, k, C)
480        let (n_rows, n_cols, k) = self;
481        let device = DeviceCpu::default();
482        let default_order = device.default_order();
483        eye_f((n_rows, n_cols, k, default_order, &device))
484    }
485}
486
487impl<T> EyeAPI<T> for usize
488where
489    T: Num + Clone + Send + Sync,
490{
491    type Out = Tensor<T, DeviceCpu, IxD>;
492
493    fn eye_f(self) -> Result<Self::Out> {
494        // n_rows -> (n_rows, n_rows, 0, C)
495        let device = DeviceCpu::default();
496        let default_order = device.default_order();
497        eye_f((self, self, 0, default_order, &device))
498    }
499}
500
501/* #endregion */
502
503/* #region full */
504
505pub trait FullAPI<Inp> {
506    type Out;
507
508    fn full_f(self) -> Result<Self::Out>;
509
510    fn full(self) -> Self::Out
511    where
512        Self: Sized,
513    {
514        Self::full_f(self).rstsr_unwrap()
515    }
516}
517
518/// New tensor having a specified shape and filled with given value.
519///
520/// # See also
521///
522/// - [Python array API standard: `full`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.full.html)
523pub fn full<Args, Inp>(param: Args) -> Args::Out
524where
525    Args: FullAPI<Inp>,
526{
527    return FullAPI::full(param);
528}
529
530pub fn full_f<Args, Inp>(param: Args) -> Result<Args::Out>
531where
532    Args: FullAPI<Inp>,
533{
534    return FullAPI::full_f(param);
535}
536
537impl<T, D, B> FullAPI<(T, D)> for (Layout<D>, T, &B)
538where
539    T: Clone,
540    D: DimAPI,
541    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
542{
543    type Out = Tensor<T, B, IxD>;
544
545    fn full_f(self) -> Result<Self::Out> {
546        let (layout, fill, device) = self;
547        let idx_max = layout.size();
548        let storage = device.full_impl(idx_max, fill)?;
549        unsafe { Ok(Tensor::new_unchecked(storage, layout.into_dim()?)) }
550    }
551}
552
553impl<T, D, B> FullAPI<(T, D)> for (D, T, FlagOrder, &B)
554where
555    T: Clone,
556    D: DimAPI,
557    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
558{
559    type Out = Tensor<T, B, IxD>;
560
561    fn full_f(self) -> Result<Self::Out> {
562        let (shape, fill, order, device) = self;
563        let layout = shape.new_contig(None, order);
564        full_f((layout, fill, device))
565    }
566}
567
568impl<T, D, B> FullAPI<(T, D)> for (D, T, &B)
569where
570    T: Clone,
571    D: DimAPI,
572    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
573{
574    type Out = Tensor<T, B, IxD>;
575
576    fn full_f(self) -> Result<Self::Out> {
577        let (shape, fill, device) = self;
578        let default_order = device.default_order();
579        let layout = shape.new_contig(None, default_order);
580        full_f((layout, fill, device))
581    }
582}
583
584impl<T, D> FullAPI<(T, D)> for (D, T, FlagOrder)
585where
586    T: Clone,
587    D: DimAPI,
588{
589    type Out = Tensor<T, DeviceCpu, IxD>;
590
591    fn full_f(self) -> Result<Self::Out> {
592        let (shape, fill, order) = self;
593        full_f((shape, fill, order, &DeviceCpu::default()))
594    }
595}
596
597#[duplicate_item(L; [D]; [Layout<D>])]
598impl<T, D> FullAPI<(T, D)> for (L, T)
599where
600    T: Clone,
601    D: DimAPI,
602{
603    type Out = Tensor<T, DeviceCpu, IxD>;
604
605    fn full_f(self) -> Result<Self::Out> {
606        let (shape, fill) = self;
607        full_f((shape, fill, &DeviceCpu::default()))
608    }
609}
610
611/* #endregion */
612
613/* #region full_like */
614
615pub trait FullLikeAPI<Inp> {
616    type Out;
617
618    fn full_like_f(self) -> Result<Self::Out>;
619
620    fn full_like(self) -> Self::Out
621    where
622        Self: Sized,
623    {
624        Self::full_like_f(self).rstsr_unwrap()
625    }
626}
627
628/// New tensor filled with given value and having the same shape as an input
629/// tensor.
630///
631/// # See also
632///
633/// - [Python array API standard: `full_like`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.full_like.html)
634pub fn full_like<Args, Inp>(param: Args) -> Args::Out
635where
636    Args: FullLikeAPI<Inp>,
637{
638    return FullLikeAPI::full_like(param);
639}
640
641pub fn full_like_f<Args, Inp>(param: Args) -> Result<Args::Out>
642where
643    Args: FullLikeAPI<Inp>,
644{
645    return FullLikeAPI::full_like_f(param);
646}
647
648impl<R, T, B, D> FullLikeAPI<()> for (&TensorAny<R, T, B, D>, T, TensorIterOrder, &B)
649where
650    T: Clone,
651    R: DataAPI<Data = B::Raw>,
652    D: DimAPI,
653    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
654{
655    type Out = Tensor<T, B, D>;
656
657    fn full_like_f(self) -> Result<Self::Out> {
658        let (tensor, fill, order, device) = self;
659        let layout = layout_for_array_copy(tensor.layout(), order)?;
660        let idx_max = layout.size();
661        let storage = device.full_impl(idx_max, fill)?;
662        unsafe { Ok(Tensor::new_unchecked(storage, layout)) }
663    }
664}
665
666impl<R, T, B, D> FullLikeAPI<()> for (&TensorAny<R, T, B, D>, T, &B)
667where
668    T: Clone,
669    R: DataAPI<Data = B::Raw>,
670    D: DimAPI,
671    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
672{
673    type Out = Tensor<T, B, D>;
674
675    fn full_like_f(self) -> Result<Self::Out> {
676        let (tensor, fill, device) = self;
677        full_like_f((tensor, fill, TensorIterOrder::default(), device))
678    }
679}
680
681impl<R, T, B, D> FullLikeAPI<()> for (&TensorAny<R, T, B, D>, T, TensorIterOrder)
682where
683    T: Clone,
684    R: DataAPI<Data = B::Raw>,
685    D: DimAPI,
686    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
687{
688    type Out = Tensor<T, B, D>;
689
690    fn full_like_f(self) -> Result<Self::Out> {
691        let (tensor, fill, order) = self;
692        let device = tensor.device();
693        full_like_f((tensor, fill, order, device))
694    }
695}
696
697impl<R, T, B, D> FullLikeAPI<()> for (&TensorAny<R, T, B, D>, T)
698where
699    T: Clone,
700    R: DataAPI<Data = B::Raw>,
701    D: DimAPI,
702    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
703{
704    type Out = Tensor<T, B, D>;
705
706    fn full_like_f(self) -> Result<Self::Out> {
707        let (tensor, fill) = self;
708        let device = tensor.device();
709        full_like_f((tensor, fill, TensorIterOrder::default(), device))
710    }
711}
712
713impl<R, T, B, D> TensorAny<R, T, B, D>
714where
715    R: DataAPI<Data = B::Raw>,
716    T: Clone,
717    D: DimAPI,
718    B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
719{
720    pub fn full_like(&self, fill: T) -> Tensor<T, B, D> {
721        full_like((self, fill))
722    }
723
724    pub fn full_like_f(&self, fill: T) -> Result<Tensor<T, B, D>> {
725        full_like_f((self, fill))
726    }
727}
728
729/* #endregion */
730
731/* #region linspace */
732
733pub trait LinspaceAPI<Inp> {
734    type Out;
735
736    fn linspace_f(self) -> Result<Self::Out>;
737
738    fn linspace(self) -> Self::Out
739    where
740        Self: Sized,
741    {
742        Self::linspace_f(self).rstsr_unwrap()
743    }
744}
745
746/// Evenly spaced numbers over a specified interval.
747///
748/// For boundary condition, current implementation is similar to numpy,
749/// where `n = 0` will return an empty array, and `n = 1` will return an
750/// array with starting value.
751///
752/// # See also
753///
754/// - [Python array API standard: `linspace`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.linspace.html)
755pub fn linspace<Args, Inp>(param: Args) -> Args::Out
756where
757    Args: LinspaceAPI<Inp>,
758{
759    return LinspaceAPI::linspace(param);
760}
761
762pub fn linspace_f<Args, Inp>(param: Args) -> Result<Args::Out>
763where
764    Args: LinspaceAPI<Inp>,
765{
766    return LinspaceAPI::linspace_f(param);
767}
768
769impl<T, B> LinspaceAPI<(T, B)> for (T, T, usize, bool, &B)
770where
771    T: ComplexFloat,
772    B: DeviceAPI<T> + DeviceCreationComplexFloatAPI<T>,
773{
774    type Out = Tensor<T, B, IxD>;
775
776    fn linspace_f(self) -> Result<Self::Out> {
777        let (start, end, n, endpoint, device) = self;
778        let data = B::linspace_impl(device, start, end, n, endpoint)?;
779        let layout = vec![data.len()].into();
780        unsafe { Ok(Tensor::new_unchecked(data, layout)) }
781    }
782}
783
784impl<T, B> LinspaceAPI<(T, B)> for (T, T, usize, &B)
785where
786    T: ComplexFloat,
787    B: DeviceAPI<T> + DeviceCreationComplexFloatAPI<T>,
788{
789    type Out = Tensor<T, B, IxD>;
790
791    fn linspace_f(self) -> Result<Self::Out> {
792        // (start, end, n, device) -> (start, end, n, true, device)
793        let (start, end, n, device) = self;
794        linspace_f((start, end, n, true, device))
795    }
796}
797
798impl<T> LinspaceAPI<T> for (T, T, usize, bool)
799where
800    T: ComplexFloat + Send + Sync,
801{
802    type Out = Tensor<T, DeviceCpu, IxD>;
803
804    fn linspace_f(self) -> Result<Self::Out> {
805        // (start, end, n, endpoint) -> (start, end, n, endpoint, device)
806        let (start, end, n, endpoint) = self;
807        linspace_f((start, end, n, endpoint, &DeviceCpu::default()))
808    }
809}
810
811impl<T> LinspaceAPI<T> for (T, T, usize)
812where
813    T: ComplexFloat + Send + Sync,
814{
815    type Out = Tensor<T, DeviceCpu, IxD>;
816
817    fn linspace_f(self) -> Result<Self::Out> {
818        // (start, end, n) -> (start, end, n, true, device)
819        let (start, end, n) = self;
820        linspace_f((start, end, n, true, &DeviceCpu::default()))
821    }
822}
823
824/* #endregion */
825
826/* #region ones */
827
828pub trait OnesAPI<Inp> {
829    type Out;
830
831    fn ones_f(self) -> Result<Self::Out>;
832
833    fn ones(self) -> Self::Out
834    where
835        Self: Sized,
836    {
837        Self::ones_f(self).rstsr_unwrap()
838    }
839}
840
841/// New tensor filled with ones and having a specified shape.
842///
843/// # See also
844///
845/// - [Python array API standard: `ones`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.ones.html)
846pub fn ones<Args, Inp>(param: Args) -> Args::Out
847where
848    Args: OnesAPI<Inp>,
849{
850    return OnesAPI::ones(param);
851}
852
853pub fn ones_f<Args, Inp>(param: Args) -> Result<Args::Out>
854where
855    Args: OnesAPI<Inp>,
856{
857    return OnesAPI::ones_f(param);
858}
859
860impl<T, D, B> OnesAPI<(T, D)> for (Layout<D>, &B)
861where
862    T: Num,
863    D: DimAPI,
864    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
865{
866    type Out = Tensor<T, B, IxD>;
867
868    fn ones_f(self) -> Result<Self::Out> {
869        let (layout, device) = self;
870        let (_, idx_max) = layout.bounds_index()?;
871        let storage = device.ones_impl(idx_max)?;
872        unsafe { Ok(Tensor::new_unchecked(storage, layout.into_dim()?)) }
873    }
874}
875
876impl<T, D, B> OnesAPI<(T, D)> for (D, FlagOrder, &B)
877where
878    T: Num,
879    D: DimAPI,
880    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
881{
882    type Out = Tensor<T, B, IxD>;
883
884    fn ones_f(self) -> Result<Self::Out> {
885        let (shape, order, device) = self;
886        let layout = shape.new_contig(None, order);
887        ones_f((layout, device))
888    }
889}
890
891impl<T, D, B> OnesAPI<(T, D)> for (D, &B)
892where
893    T: Num,
894    D: DimAPI,
895    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
896{
897    type Out = Tensor<T, B, IxD>;
898
899    fn ones_f(self) -> Result<Self::Out> {
900        let (shape, device) = self;
901        let default_order = device.default_order();
902        let layout = shape.new_contig(None, default_order);
903        ones_f((layout, device))
904    }
905}
906
907impl<T, D> OnesAPI<(T, D)> for (D, FlagOrder)
908where
909    T: Num + Clone,
910    D: DimAPI,
911{
912    type Out = Tensor<T, DeviceCpu, IxD>;
913
914    fn ones_f(self) -> Result<Self::Out> {
915        let (shape, order) = self;
916        ones_f((shape, order, &DeviceCpu::default()))
917    }
918}
919
920#[duplicate_item(L; [D]; [Layout<D>])]
921impl<T, D> OnesAPI<(T, D)> for L
922where
923    T: Num + Clone,
924    D: DimAPI,
925{
926    type Out = Tensor<T, DeviceCpu, IxD>;
927
928    fn ones_f(self) -> Result<Self::Out> {
929        ones_f((self, &DeviceCpu::default()))
930    }
931}
932
933/* #endregion */
934
935/* #region ones_like */
936
937pub trait OnesLikeAPI<Inp> {
938    type Out;
939
940    fn ones_like_f(self) -> Result<Self::Out>;
941
942    fn ones_like(self) -> Self::Out
943    where
944        Self: Sized,
945    {
946        Self::ones_like_f(self).rstsr_unwrap()
947    }
948}
949
950/// New tensor filled with ones and having the same shape as an input
951/// tensor.
952///
953/// # See also
954///
955/// - [Python array API standard: `ones_like`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.ones_like.html)
956pub fn ones_like<Args, Inp>(param: Args) -> Args::Out
957where
958    Args: OnesLikeAPI<Inp>,
959{
960    return OnesLikeAPI::ones_like(param);
961}
962
963pub fn ones_like_f<Args, Inp>(param: Args) -> Result<Args::Out>
964where
965    Args: OnesLikeAPI<Inp>,
966{
967    return OnesLikeAPI::ones_like_f(param);
968}
969
970impl<R, T, B, D> OnesLikeAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder, &B)
971where
972    R: DataAPI<Data = B::Raw>,
973    T: Num,
974    D: DimAPI,
975    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
976{
977    type Out = Tensor<T, B, D>;
978
979    fn ones_like_f(self) -> Result<Self::Out> {
980        let (tensor, order, device) = self;
981        let layout = layout_for_array_copy(tensor.layout(), order)?;
982        let idx_max = layout.size();
983        let storage = device.ones_impl(idx_max)?;
984        unsafe { Ok(Tensor::new_unchecked(storage, layout)) }
985    }
986}
987
988impl<R, T, B, D> OnesLikeAPI<()> for (&TensorAny<R, T, B, D>, &B)
989where
990    R: DataAPI<Data = B::Raw>,
991    T: Num,
992    D: DimAPI,
993    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
994{
995    type Out = Tensor<T, B, D>;
996
997    fn ones_like_f(self) -> Result<Self::Out> {
998        let (tensor, device) = self;
999        ones_like_f((tensor, TensorIterOrder::default(), device))
1000    }
1001}
1002
1003impl<R, T, B, D> OnesLikeAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder)
1004where
1005    R: DataAPI<Data = B::Raw>,
1006    T: Num,
1007    D: DimAPI,
1008    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1009{
1010    type Out = Tensor<T, B, D>;
1011
1012    fn ones_like_f(self) -> Result<Self::Out> {
1013        let (tensor, order) = self;
1014        let device = tensor.device();
1015        ones_like_f((tensor, order, device))
1016    }
1017}
1018
1019impl<R, T, B, D> OnesLikeAPI<()> for &TensorAny<R, T, B, D>
1020where
1021    R: DataAPI<Data = B::Raw>,
1022    T: Num,
1023    D: DimAPI,
1024    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1025{
1026    type Out = Tensor<T, B, D>;
1027
1028    fn ones_like_f(self) -> Result<Self::Out> {
1029        let device = self.device();
1030        ones_like_f((self, TensorIterOrder::default(), device))
1031    }
1032}
1033
1034impl<R, T, B, D> TensorAny<R, T, B, D>
1035where
1036    R: DataAPI<Data = B::Raw>,
1037    D: DimAPI,
1038    T: Num,
1039    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1040{
1041    /// New tensor filled with ones and having the same shape as an input
1042    /// tensor.
1043    ///
1044    /// # See also
1045    ///
1046    /// [`ones_like`]
1047    pub fn ones_like(&self) -> Tensor<T, B, D> {
1048        ones_like((self, TensorIterOrder::default(), self.device()))
1049    }
1050
1051    pub fn ones_like_f(&self) -> Result<Tensor<T, B, D>> {
1052        ones_like_f((self, TensorIterOrder::default(), self.device()))
1053    }
1054}
1055
1056/* #endregion */
1057
1058/* #region uninit */
1059
1060pub trait UninitAPI<Inp> {
1061    type Out;
1062
1063    fn uninit_f(self) -> Result<Self::Out>;
1064
1065    fn uninit(self) -> Self::Out
1066    where
1067        Self: Sized,
1068    {
1069        Self::uninit_f(self).rstsr_unwrap()
1070    }
1071}
1072
1073/// New tensor filled with uninitialized values and having a specified shape.
1074pub fn uninit<Args, Inp>(param: Args) -> Args::Out
1075where
1076    Args: UninitAPI<Inp>,
1077{
1078    return UninitAPI::uninit(param);
1079}
1080
1081pub fn uninit_f<Args, Inp>(param: Args) -> Result<Args::Out>
1082where
1083    Args: UninitAPI<Inp>,
1084{
1085    return UninitAPI::uninit_f(param);
1086}
1087
1088impl<T, D, B> UninitAPI<(T, D)> for (Layout<D>, &B)
1089where
1090    D: DimAPI,
1091    B: DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T>,
1092{
1093    type Out = Tensor<MaybeUninit<T>, B, IxD>;
1094
1095    fn uninit_f(self) -> Result<Self::Out> {
1096        let (layout, device) = self;
1097        let (_, idx_max) = layout.bounds_index()?;
1098        let storage = B::uninit_impl(device, idx_max)?;
1099        unsafe { Ok(Tensor::new_unchecked(storage, layout.into_dim()?)) }
1100    }
1101}
1102
1103impl<T, D, B> UninitAPI<(T, D)> for (D, FlagOrder, &B)
1104where
1105    D: DimAPI,
1106    B: DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T>,
1107{
1108    type Out = Tensor<MaybeUninit<T>, B, IxD>;
1109
1110    fn uninit_f(self) -> Result<Self::Out> {
1111        let (shape, order, device) = self;
1112        let layout = shape.new_contig(None, order);
1113        uninit_f((layout, device))
1114    }
1115}
1116
1117impl<T, D, B> UninitAPI<(T, D)> for (D, &B)
1118where
1119    D: DimAPI,
1120    B: DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T>,
1121{
1122    type Out = Tensor<MaybeUninit<T>, B, IxD>;
1123
1124    fn uninit_f(self) -> Result<Self::Out> {
1125        let (shape, device) = self;
1126        let default_order = device.default_order();
1127        let layout = shape.new_contig(None, default_order);
1128        uninit_f((layout, device))
1129    }
1130}
1131
1132impl<T, D> UninitAPI<(T, D)> for (D, FlagOrder)
1133where
1134    T: Clone,
1135    D: DimAPI,
1136{
1137    type Out = Tensor<MaybeUninit<T>, DeviceCpu, IxD>;
1138
1139    fn uninit_f(self) -> Result<Self::Out> {
1140        let (shape, order) = self;
1141        uninit_f((shape, order, &DeviceCpu::default()))
1142    }
1143}
1144
1145#[duplicate_item(L; [D]; [Layout<D>])]
1146impl<T, D> UninitAPI<(T, D)> for L
1147where
1148    T: Clone,
1149    D: DimAPI,
1150{
1151    type Out = Tensor<MaybeUninit<T>, DeviceCpu, IxD>;
1152
1153    fn uninit_f(self) -> Result<Self::Out> {
1154        uninit_f((self, &DeviceCpu::default()))
1155    }
1156}
1157
1158/* #endregion */
1159
1160/* #region assume_init */
1161
1162/// Converts a tensor with uninitialized values into a tensor with initialized values.
1163///
1164/// # Safety
1165///
1166/// This function is unsafe because it assumes that all elements in the input tensor are properly
1167/// initialized.
1168pub unsafe fn assume_init_f<T, B, D>(tensor: Tensor<MaybeUninit<T>, B, D>) -> Result<Tensor<T, B, D>>
1169where
1170    D: DimAPI,
1171    B: DeviceAPI<T> + DeviceAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T>,
1172{
1173    let (storage, layout) = tensor.into_raw_parts();
1174    let storage = B::assume_init_impl(storage)?;
1175    unsafe { Ok(Tensor::new_unchecked(storage, layout)) }
1176}
1177
1178/// Converts a tensor with uninitialized values into a tensor with initialized values.
1179///
1180/// # Safety
1181///
1182/// This function is unsafe because it assumes that all elements in the input tensor are properly
1183/// initialized.
1184pub unsafe fn assume_init<T, B, D>(tensor: Tensor<MaybeUninit<T>, B, D>) -> Tensor<T, B, D>
1185where
1186    D: DimAPI,
1187    B: DeviceAPI<T> + DeviceAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T>,
1188{
1189    unsafe { assume_init_f(tensor).rstsr_unwrap() }
1190}
1191
1192/* #endregion */
1193
1194/* #region zeros */
1195
1196pub trait ZerosAPI<Inp> {
1197    type Out;
1198
1199    fn zeros_f(self) -> Result<Self::Out>;
1200
1201    fn zeros(self) -> Self::Out
1202    where
1203        Self: Sized,
1204    {
1205        Self::zeros_f(self).rstsr_unwrap()
1206    }
1207}
1208
1209/// New tensor filled with zeros and having a specified shape.
1210///
1211/// # See also
1212///
1213/// - [Python array API standard: `zeros`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.zeros.html)
1214pub fn zeros<Args, Inp>(param: Args) -> Args::Out
1215where
1216    Args: ZerosAPI<Inp>,
1217{
1218    return ZerosAPI::zeros(param);
1219}
1220
1221pub fn zeros_f<Args, Inp>(param: Args) -> Result<Args::Out>
1222where
1223    Args: ZerosAPI<Inp>,
1224{
1225    return ZerosAPI::zeros_f(param);
1226}
1227
1228impl<T, D, B> ZerosAPI<(T, D)> for (Layout<D>, &B)
1229where
1230    T: Num,
1231    D: DimAPI,
1232    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1233{
1234    type Out = Tensor<T, B, IxD>;
1235
1236    fn zeros_f(self) -> Result<Self::Out> {
1237        let (layout, device) = self;
1238        let (_, idx_max) = layout.bounds_index()?;
1239        let storage = B::zeros_impl(device, idx_max)?;
1240        unsafe { Ok(Tensor::new_unchecked(storage, layout.into_dim()?)) }
1241    }
1242}
1243
1244impl<T, D, B> ZerosAPI<(T, D)> for (D, FlagOrder, &B)
1245where
1246    T: Num,
1247    D: DimAPI,
1248    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1249{
1250    type Out = Tensor<T, B, IxD>;
1251
1252    fn zeros_f(self) -> Result<Self::Out> {
1253        let (shape, order, device) = self;
1254        let layout = shape.new_contig(None, order);
1255        zeros_f((layout, device))
1256    }
1257}
1258
1259impl<T, D, B> ZerosAPI<(T, D)> for (D, &B)
1260where
1261    T: Num,
1262    D: DimAPI,
1263    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1264{
1265    type Out = Tensor<T, B, IxD>;
1266
1267    fn zeros_f(self) -> Result<Self::Out> {
1268        let (shape, device) = self;
1269        let default_order = device.default_order();
1270        let layout = shape.new_contig(None, default_order);
1271        zeros_f((layout, device))
1272    }
1273}
1274
1275impl<T, D> ZerosAPI<(T, D)> for (D, FlagOrder)
1276where
1277    T: Num + Clone,
1278    D: DimAPI,
1279{
1280    type Out = Tensor<T, DeviceCpu, IxD>;
1281
1282    fn zeros_f(self) -> Result<Self::Out> {
1283        let (shape, order) = self;
1284        zeros_f((shape, order, &DeviceCpu::default()))
1285    }
1286}
1287
1288#[duplicate_item(L; [D]; [Layout<D>])]
1289impl<T, D> ZerosAPI<(T, D)> for L
1290where
1291    T: Num + Clone,
1292    D: DimAPI,
1293{
1294    type Out = Tensor<T, DeviceCpu, IxD>;
1295
1296    fn zeros_f(self) -> Result<Self::Out> {
1297        zeros_f((self, &DeviceCpu::default()))
1298    }
1299}
1300
1301/* #endregion */
1302
1303/* #region zeros_like */
1304
1305pub trait ZerosLikeAPI<Inp> {
1306    type Out;
1307
1308    fn zeros_like_f(self) -> Result<Self::Out>;
1309
1310    fn zeros_like(self) -> Self::Out
1311    where
1312        Self: Sized,
1313    {
1314        Self::zeros_like_f(self).rstsr_unwrap()
1315    }
1316}
1317
1318/// New tensor filled with zeros and having the same shape as an input
1319/// tensor.
1320///
1321/// # See also
1322///
1323/// - [Python array API standard: `zeros_like`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.zeros_like.html)
1324pub fn zeros_like<Args, Inp>(param: Args) -> Args::Out
1325where
1326    Args: ZerosLikeAPI<Inp>,
1327{
1328    return ZerosLikeAPI::zeros_like(param);
1329}
1330
1331pub fn zeros_like_f<Args, Inp>(param: Args) -> Result<Args::Out>
1332where
1333    Args: ZerosLikeAPI<Inp>,
1334{
1335    return ZerosLikeAPI::zeros_like_f(param);
1336}
1337
1338impl<R, T, B, D> ZerosLikeAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder, &B)
1339where
1340    R: DataAPI<Data = B::Raw>,
1341    T: Num,
1342    D: DimAPI,
1343    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1344{
1345    type Out = Tensor<T, B, D>;
1346
1347    fn zeros_like_f(self) -> Result<Self::Out> {
1348        let (tensor, order, device) = self;
1349        let layout = layout_for_array_copy(tensor.layout(), order)?;
1350        let idx_max = layout.size();
1351        let storage = B::zeros_impl(device, idx_max)?;
1352        unsafe { Ok(Tensor::new_unchecked(storage, layout)) }
1353    }
1354}
1355
1356impl<R, T, B, D> ZerosLikeAPI<()> for (&TensorAny<R, T, B, D>, &B)
1357where
1358    R: DataAPI<Data = B::Raw>,
1359    T: Num,
1360    D: DimAPI,
1361    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1362{
1363    type Out = Tensor<T, B, D>;
1364
1365    fn zeros_like_f(self) -> Result<Self::Out> {
1366        let (tensor, device) = self;
1367        zeros_like_f((tensor, TensorIterOrder::default(), device))
1368    }
1369}
1370
1371impl<R, T, B, D> ZerosLikeAPI<()> for (&TensorAny<R, T, B, D>, TensorIterOrder)
1372where
1373    R: DataAPI<Data = B::Raw>,
1374    T: Num,
1375    D: DimAPI,
1376    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1377{
1378    type Out = Tensor<T, B, D>;
1379
1380    fn zeros_like_f(self) -> Result<Self::Out> {
1381        let (tensor, order) = self;
1382        let device = tensor.device();
1383        zeros_like_f((tensor, order, device))
1384    }
1385}
1386
1387impl<R, T, B, D> ZerosLikeAPI<()> for &TensorAny<R, T, B, D>
1388where
1389    R: DataAPI<Data = B::Raw>,
1390    T: Num,
1391    D: DimAPI,
1392    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1393{
1394    type Out = Tensor<T, B, D>;
1395
1396    fn zeros_like_f(self) -> Result<Self::Out> {
1397        zeros_like_f((self, TensorIterOrder::default(), self.device()))
1398    }
1399}
1400
1401impl<R, T, B, D> TensorAny<R, T, B, D>
1402where
1403    R: DataAPI<Data = B::Raw>,
1404    D: DimAPI,
1405    T: Num,
1406    B: DeviceAPI<T> + DeviceCreationNumAPI<T>,
1407{
1408    /// New tensor filled with zeros and having the same shape as an input
1409    /// tensor.
1410    ///
1411    /// # See also
1412    ///
1413    /// [`zeros_like`]
1414    pub fn zeros_like(&self) -> Tensor<T, B, D> {
1415        zeros_like((self, TensorIterOrder::default(), self.device()))
1416    }
1417
1418    pub fn zeros_like_f(&self) -> Result<Tensor<T, B, D>> {
1419        zeros_like_f((self, TensorIterOrder::default(), self.device()))
1420    }
1421}
1422
1423/* #endregion */
1424
1425/* #region tril */
1426
1427pub trait TrilAPI<Inp> {
1428    type Out;
1429
1430    fn tril_f(self) -> Result<Self::Out>;
1431
1432    fn tril(self) -> Self::Out
1433    where
1434        Self: Sized,
1435    {
1436        Self::tril_f(self).rstsr_unwrap()
1437    }
1438}
1439
1440/// Returns the lower triangular part of a matrix (or a stack of matrices) x.
1441///
1442/// # See also
1443///
1444/// - [Python array API standard: `tril`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.tril.html)
1445pub fn tril<Args, Inp>(param: Args) -> Args::Out
1446where
1447    Args: TrilAPI<Inp>,
1448{
1449    return TrilAPI::tril(param);
1450}
1451
1452pub fn tril_f<Args, Inp>(param: Args) -> Result<Args::Out>
1453where
1454    Args: TrilAPI<Inp>,
1455{
1456    return TrilAPI::tril_f(param);
1457}
1458
1459impl<T, D, B> TrilAPI<()> for (TensorView<'_, T, B, D>, isize)
1460where
1461    T: Num + Clone,
1462    D: DimAPI,
1463    B: DeviceAPI<T>
1464        + DeviceRawAPI<MaybeUninit<T>>
1465        + DeviceCreationTriAPI<T>
1466        + DeviceCreationAnyAPI<T>
1467        + OpAssignArbitaryAPI<T, D, D>
1468        + OpAssignAPI<T, D>,
1469    <B as DeviceRawAPI<T>>::Raw: Clone,
1470{
1471    type Out = Tensor<T, B, D>;
1472
1473    fn tril_f(self) -> Result<Self::Out> {
1474        let (x, k) = self;
1475        let default_order = x.device().default_order();
1476        let mut x = x.into_contig_f(default_order)?;
1477        let device = x.device().clone();
1478        let layout = x.layout().clone();
1479        device.tril_impl(x.raw_mut(), &layout, k)?;
1480        Ok(x)
1481    }
1482}
1483
1484impl<T, D, B> TrilAPI<()> for TensorView<'_, T, B, D>
1485where
1486    T: Num + Clone,
1487    D: DimAPI,
1488    B: DeviceAPI<T>
1489        + DeviceRawAPI<MaybeUninit<T>>
1490        + DeviceCreationTriAPI<T>
1491        + DeviceCreationAnyAPI<T>
1492        + OpAssignArbitaryAPI<T, D, D>
1493        + OpAssignAPI<T, D>,
1494    <B as DeviceRawAPI<T>>::Raw: Clone,
1495{
1496    type Out = Tensor<T, B, D>;
1497
1498    fn tril_f(self) -> Result<Self::Out> {
1499        tril_f((self, 0))
1500    }
1501}
1502
1503impl<'a, T, D, B> TrilAPI<()> for (TensorMut<'a, T, B, D>, isize)
1504where
1505    T: Num + Clone,
1506    D: DimAPI,
1507    B: DeviceAPI<T>
1508        + DeviceCreationTriAPI<T>
1509        + DeviceCreationAnyAPI<T>
1510        + OpAssignArbitaryAPI<T, D, D>
1511        + OpAssignAPI<T, D>,
1512{
1513    type Out = TensorMut<'a, T, B, D>;
1514
1515    fn tril_f(self) -> Result<Self::Out> {
1516        let (mut x, k) = self;
1517        let device = x.device().clone();
1518        let layout = x.layout().clone();
1519        device.tril_impl(x.raw_mut(), &layout, k)?;
1520        Ok(x)
1521    }
1522}
1523
1524impl<'a, T, D, B> TrilAPI<()> for TensorMut<'a, T, B, D>
1525where
1526    T: Num + Clone,
1527    D: DimAPI,
1528    B: DeviceAPI<T>
1529        + DeviceCreationTriAPI<T>
1530        + DeviceCreationAnyAPI<T>
1531        + OpAssignArbitaryAPI<T, D, D>
1532        + OpAssignAPI<T, D>,
1533{
1534    type Out = TensorMut<'a, T, B, D>;
1535
1536    fn tril_f(self) -> Result<Self::Out> {
1537        tril_f((self, 0))
1538    }
1539}
1540
1541impl<T, D, B> TrilAPI<()> for (Tensor<T, B, D>, isize)
1542where
1543    T: Num + Clone,
1544    D: DimAPI,
1545    B: DeviceAPI<T> + DeviceCreationTriAPI<T> + DeviceCreationAnyAPI<T>,
1546{
1547    type Out = Tensor<T, B, D>;
1548
1549    fn tril_f(self) -> Result<Self::Out> {
1550        let (mut x, k) = self;
1551        let device = x.device().clone();
1552        let layout = x.layout().clone();
1553        device.tril_impl(x.raw_mut(), &layout, k)?;
1554        Ok(x)
1555    }
1556}
1557
1558impl<T, D, B> TrilAPI<()> for Tensor<T, B, D>
1559where
1560    T: Num + Clone,
1561    D: DimAPI,
1562    B: DeviceAPI<T> + DeviceCreationTriAPI<T> + DeviceCreationAnyAPI<T>,
1563{
1564    type Out = Tensor<T, B, D>;
1565
1566    fn tril_f(self) -> Result<Self::Out> {
1567        tril_f((self, 0))
1568    }
1569}
1570
1571impl<R, T, D, B> TrilAPI<()> for (&TensorAny<R, T, B, D>, isize)
1572where
1573    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
1574    T: Num + Clone,
1575    D: DimAPI,
1576    B: DeviceAPI<T>
1577        + DeviceRawAPI<MaybeUninit<T>>
1578        + DeviceCreationTriAPI<T>
1579        + DeviceCreationAnyAPI<T>
1580        + OpAssignArbitaryAPI<T, D, D>
1581        + OpAssignAPI<T, D>,
1582    <B as DeviceRawAPI<T>>::Raw: Clone,
1583{
1584    type Out = Tensor<T, B, D>;
1585
1586    fn tril_f(self) -> Result<Self::Out> {
1587        let (x, k) = self;
1588        tril_f((x.view(), k))
1589    }
1590}
1591
1592impl<R, T, D, B> TrilAPI<()> for &TensorAny<R, T, B, D>
1593where
1594    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
1595    T: Num + Clone,
1596    D: DimAPI,
1597    B: DeviceAPI<T>
1598        + DeviceRawAPI<MaybeUninit<T>>
1599        + DeviceCreationTriAPI<T>
1600        + DeviceCreationAnyAPI<T>
1601        + OpAssignArbitaryAPI<T, D, D>
1602        + OpAssignAPI<T, D>,
1603    <B as DeviceRawAPI<T>>::Raw: Clone,
1604{
1605    type Out = Tensor<T, B, D>;
1606
1607    fn tril_f(self) -> Result<Self::Out> {
1608        tril_f((self.view(), 0))
1609    }
1610}
1611
1612/* #endregion */
1613
1614/* #region triu */
1615
1616pub trait TriuAPI<Inp> {
1617    type Out;
1618
1619    fn triu_f(self) -> Result<Self::Out>;
1620
1621    fn triu(self) -> Self::Out
1622    where
1623        Self: Sized,
1624    {
1625        Self::triu_f(self).rstsr_unwrap()
1626    }
1627}
1628
1629/// Returns the upper triangular part of a matrix (or a stack of matrices) x.
1630///
1631/// # See also
1632///
1633/// - [Python array API standard: `triu`](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.triu.html)
1634pub fn triu<Args, Inp>(param: Args) -> Args::Out
1635where
1636    Args: TriuAPI<Inp>,
1637{
1638    return TriuAPI::triu(param);
1639}
1640
1641pub fn triu_f<Args, Inp>(param: Args) -> Result<Args::Out>
1642where
1643    Args: TriuAPI<Inp>,
1644{
1645    return TriuAPI::triu_f(param);
1646}
1647
1648impl<T, D, B> TriuAPI<()> for (TensorView<'_, T, B, D>, isize)
1649where
1650    T: Num + Clone,
1651    D: DimAPI,
1652    B: DeviceAPI<T>
1653        + DeviceRawAPI<MaybeUninit<T>>
1654        + DeviceCreationTriAPI<T>
1655        + DeviceCreationAnyAPI<T>
1656        + OpAssignArbitaryAPI<T, D, D>
1657        + OpAssignAPI<T, D>,
1658    <B as DeviceRawAPI<T>>::Raw: Clone,
1659{
1660    type Out = Tensor<T, B, D>;
1661
1662    fn triu_f(self) -> Result<Self::Out> {
1663        let (x, k) = self;
1664        let default_order = x.device().default_order();
1665        let mut x = x.into_contig_f(default_order)?;
1666        let device = x.device().clone();
1667        let layout = x.layout().clone();
1668        device.triu_impl(x.raw_mut(), &layout, k)?;
1669        Ok(x)
1670    }
1671}
1672
1673impl<T, D, B> TriuAPI<()> for TensorView<'_, T, B, D>
1674where
1675    T: Num + Clone,
1676    D: DimAPI,
1677    B: DeviceAPI<T>
1678        + DeviceRawAPI<MaybeUninit<T>>
1679        + DeviceCreationTriAPI<T>
1680        + DeviceCreationAnyAPI<T>
1681        + OpAssignArbitaryAPI<T, D, D>
1682        + OpAssignAPI<T, D>,
1683    <B as DeviceRawAPI<T>>::Raw: Clone,
1684{
1685    type Out = Tensor<T, B, D>;
1686
1687    fn triu_f(self) -> Result<Self::Out> {
1688        triu_f((self, 0))
1689    }
1690}
1691
1692impl<'a, T, D, B> TriuAPI<()> for (TensorMut<'a, T, B, D>, isize)
1693where
1694    T: Num + Clone,
1695    D: DimAPI,
1696    B: DeviceAPI<T>
1697        + DeviceCreationTriAPI<T>
1698        + DeviceCreationAnyAPI<T>
1699        + OpAssignArbitaryAPI<T, D, D>
1700        + OpAssignAPI<T, D>,
1701{
1702    type Out = TensorMut<'a, T, B, D>;
1703
1704    fn triu_f(self) -> Result<Self::Out> {
1705        let (mut x, k) = self;
1706        let device = x.device().clone();
1707        let layout = x.layout().clone();
1708        device.triu_impl(x.raw_mut(), &layout, k)?;
1709        Ok(x)
1710    }
1711}
1712
1713impl<'a, T, D, B> TriuAPI<()> for TensorMut<'a, T, B, D>
1714where
1715    T: Num + Clone,
1716    D: DimAPI,
1717    B: DeviceAPI<T>
1718        + DeviceCreationTriAPI<T>
1719        + DeviceCreationAnyAPI<T>
1720        + OpAssignArbitaryAPI<T, D, D>
1721        + OpAssignAPI<T, D>,
1722{
1723    type Out = TensorMut<'a, T, B, D>;
1724
1725    fn triu_f(self) -> Result<Self::Out> {
1726        triu_f((self, 0))
1727    }
1728}
1729
1730impl<T, D, B> TriuAPI<()> for (Tensor<T, B, D>, isize)
1731where
1732    T: Num + Clone,
1733    D: DimAPI,
1734    B: DeviceAPI<T> + DeviceCreationTriAPI<T> + DeviceCreationAnyAPI<T>,
1735{
1736    type Out = Tensor<T, B, D>;
1737
1738    fn triu_f(self) -> Result<Self::Out> {
1739        let (mut x, k) = self;
1740        let device = x.device().clone();
1741        let layout = x.layout().clone();
1742        device.triu_impl(x.raw_mut(), &layout, k)?;
1743        Ok(x)
1744    }
1745}
1746
1747impl<T, D, B> TriuAPI<()> for Tensor<T, B, D>
1748where
1749    T: Num + Clone,
1750    D: DimAPI,
1751    B: DeviceAPI<T> + DeviceCreationTriAPI<T> + DeviceCreationAnyAPI<T>,
1752{
1753    type Out = Tensor<T, B, D>;
1754
1755    fn triu_f(self) -> Result<Self::Out> {
1756        triu_f((self, 0))
1757    }
1758}
1759
1760impl<R, T, D, B> TriuAPI<()> for (&TensorAny<R, T, B, D>, isize)
1761where
1762    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
1763    T: Num + Clone,
1764    D: DimAPI,
1765    B: DeviceAPI<T>
1766        + DeviceRawAPI<MaybeUninit<T>>
1767        + DeviceCreationTriAPI<T>
1768        + DeviceCreationAnyAPI<T>
1769        + OpAssignArbitaryAPI<T, D, D>
1770        + OpAssignAPI<T, D>,
1771    <B as DeviceRawAPI<T>>::Raw: Clone,
1772{
1773    type Out = Tensor<T, B, D>;
1774
1775    fn triu_f(self) -> Result<Self::Out> {
1776        let (x, k) = self;
1777        triu_f((x.view(), k))
1778    }
1779}
1780
1781impl<R, T, D, B> TriuAPI<()> for &TensorAny<R, T, B, D>
1782where
1783    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
1784    T: Num + Clone,
1785    D: DimAPI,
1786    B: DeviceAPI<T>
1787        + DeviceRawAPI<MaybeUninit<T>>
1788        + DeviceCreationTriAPI<T>
1789        + DeviceCreationAnyAPI<T>
1790        + OpAssignArbitaryAPI<T, D, D>
1791        + OpAssignAPI<T, D>,
1792    <B as DeviceRawAPI<T>>::Raw: Clone,
1793{
1794    type Out = Tensor<T, B, D>;
1795
1796    fn triu_f(self) -> Result<Self::Out> {
1797        triu_f((self.view(), 0))
1798    }
1799}
1800
1801/* #endregion */
1802
1803/* #region macro tensor_from_nested */
1804
1805pub trait FromNestedArrayAPI<Arr, B> {
1806    fn from_nested_array(arr: Arr, device: &B) -> Self;
1807}
1808
1809/// Implementation of this macro uses the same way of crate `ndarray`:
1810/// <https://docs.rs/ndarray/latest/ndarray/macro.array.html>
1811macro_rules! impl_from_nested_array {
1812    ($arr_type:ty, $ix_type:tt, $($n:ident),+) => {
1813        impl<T, B, $(const $n: usize),+> FromNestedArrayAPI<$arr_type, B> for Tensor<T, B, $ix_type>
1814        where
1815            T: Clone,
1816            B: DeviceAPI<T> + DeviceCreationAnyAPI<T>,
1817        {
1818            #[allow(clippy::missing_transmute_annotations)]
1819            fn from_nested_array(arr: $arr_type, device: &B) -> Self {
1820                let shape: Layout<$ix_type> = [$($n),+].c();
1821                let slc = unsafe {
1822                    core::slice::from_raw_parts(arr.as_ptr() as *const T, shape.size())
1823                };
1824                let vec = slc.to_vec();
1825                asarray((vec, shape, device)).into_dim()
1826            }
1827        }
1828    }
1829}
1830
1831impl_from_nested_array!([T; N1], Ix1, N1);
1832impl_from_nested_array!([[T; N2]; N1], Ix2, N1, N2);
1833impl_from_nested_array!([[[T; N3]; N2]; N1], Ix3, N1, N2, N3);
1834impl_from_nested_array!([[[[T; N4]; N3]; N2]; N1], Ix4, N1, N2, N3, N4);
1835impl_from_nested_array!([[[[[T; N5]; N4]; N3]; N2]; N1], Ix5, N1, N2, N3, N4, N5);
1836impl_from_nested_array!([[[[[[T; N6]; N5]; N4]; N3]; N2]; N1], Ix6, N1, N2, N3, N4, N5, N6);
1837
1838/// Create a tensor from a nested array literal (for development purposes).
1839///
1840/// # Examples
1841///
1842/// ```rust
1843/// use rstsr::prelude::*;
1844///
1845/// let tsr = rt::tensor_from_nested!(
1846///     [[1, 2, 3],
1847///      [4, 5, 6]]
1848/// );
1849/// println!("{tsr:?}");
1850/// // [[ 1 2 3]
1851/// //  [ 4 5 6]]
1852/// // 2-Dim (dyn), contiguous: Cc
1853///
1854/// let tsr = rt::tensor_from_nested!(
1855///     [[[1, 2],
1856///       [3, 4]],
1857///      [[5, 6],
1858///       [7, 8]]]
1859/// );
1860/// println!("{tsr:?}");
1861/// // [[[ 1 2]
1862/// //   [ 3 4]]
1863/// //  [[ 5 6]
1864/// //   [ 7 8]]]
1865/// // 3-Dim (dyn), contiguous: Cc
1866/// ```
1867///
1868/// If you want to specify the device, you can do it like this:
1869/// ```rust
1870/// use rstsr::prelude::*;
1871///
1872/// let device = DeviceFaer::default(); // or other devices
1873/// let tsr = rt::tensor_from_nested!([[1, 2, 3], [4, 5, 6]], &device);
1874/// println!("{tsr:?}"); // you will get a tensor on DeviceFaer
1875/// ```
1876///
1877/// # Notes on usage
1878///
1879/// - This macro is used for testing or prototyping purposes. For production code, it is recommended
1880///   to use other methods to create tensors (e.g., [`asarray`]) to have better control and
1881///   performance.
1882/// - This macro only gives row-major (C-contiguous) layout tensors, even if the device's default
1883///   order is column-major.
1884/// - The macro supports up to 6 dimensions. For 7 or more dimensions, use other methods to create
1885///   tensors.
1886/// - The created tensor is on the [`DeviceCpu`] (depends on crate feature `faer_as_default`) and
1887///   has dynamic dimension type [`IxD`].
1888/// - It is adviced to use brackets `[]` to denote each dimension level, instead of parentheses
1889///   `()`.
1890///
1891/// # See also
1892///
1893/// - [`ndarray::array`](https://docs.rs/ndarray/latest/ndarray/macro.array.html) from rust crate
1894///   ndarray.
1895#[macro_export]
1896macro_rules! tensor_from_nested {
1897    ([$([$([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*], $device:expr) => {{
1898        compile_error!("Tensor of 7 dimensions or more cannot be constructed with the tensor_from_nested! macro.");
1899    }};
1900    ([$([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*], $device:expr) => {{
1901        Tensor::<_, _, Ix6>::from_nested_array([$([$([$([$([$([$($x,)*],)*],)*],)*],)*],)*], $device).into_dyn()
1902    }};
1903    ([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*], $device:expr) => {{
1904        Tensor::<_, _, Ix5>::from_nested_array([$([$([$([$([$($x,)*],)*],)*],)*],)*], $device).into_dyn()
1905    }};
1906    ([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*], $device:expr) => {{
1907        Tensor::<_, _, Ix4>::from_nested_array([$([$([$([$($x,)*],)*],)*],)*], $device).into_dyn()
1908    }};
1909    ([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*], $device:expr) => {{
1910        Tensor::<_, _, Ix3>::from_nested_array([$([$([$($x,)*],)*],)*], $device).into_dyn()
1911    }};
1912    ([$([$($x:expr),* $(,)*]),+ $(,)*], $device:expr) => {{
1913        Tensor::<_, _, Ix2>::from_nested_array([$([$($x,)*],)*], $device).into_dyn()
1914    }};
1915    ([$($x:expr),* $(,)*], $device:expr) => {{
1916        Tensor::<_, _, Ix1>::from_nested_array([$($x,)*], $device).into_dyn()
1917    }};
1918    ([$([$([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]) => {{
1919        compile_error!("Tensor of 7 dimensions or more cannot be constructed with the tensor_from_nested! macro.");
1920    }};
1921    ([$([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]) => {{
1922        let device = DeviceCpu::default();
1923        Tensor::<_, _, Ix6>::from_nested_array([$([$([$([$([$([$($x,)*],)*],)*],)*],)*],)*], &device).into_dyn()
1924    }};
1925    ([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]) => {{
1926        let device = DeviceCpu::default();
1927        Tensor::<_, _, Ix5>::from_nested_array([$([$([$([$([$($x,)*],)*],)*],)*],)*], &device).into_dyn()
1928    }};
1929    ([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]) => {{
1930        let device = DeviceCpu::default();
1931        Tensor::<_, _, Ix4>::from_nested_array([$([$([$([$($x,)*],)*],)*],)*], &device).into_dyn()
1932    }};
1933    ([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]) => {{
1934        let device = DeviceCpu::default();
1935        Tensor::<_, _, Ix3>::from_nested_array([$([$([$($x,)*],)*],)*], &device).into_dyn()
1936    }};
1937    ([$([$($x:expr),* $(,)*]),+ $(,)*]) => {{
1938        let device = DeviceCpu::default();
1939        Tensor::<_, _, Ix2>::from_nested_array([$([$($x,)*],)*], &device).into_dyn()
1940    }};
1941    ([$($x:expr),* $(,)*]) => {{
1942        let device = DeviceCpu::default();
1943        Tensor::<_, _, Ix1>::from_nested_array([$($x,)*], &device).into_dyn()
1944    }};
1945}
1946
1947#[test]
1948fn playground() {
1949    use rstsr::prelude::*;
1950    let tsr = rt::tensor_from_nested!([[1, 2, 3], [4, 5, 6]]);
1951    println!("{tsr:?}");
1952
1953    let mut device = DeviceCpuSerial::default();
1954    device.set_default_order(ColMajor);
1955    let tsr = rt::tensor_from_nested!([[1, 2, 3], [4, 5, 6]], &device);
1956    println!("{tsr:?}");
1957}
1958
1959/* #endregion */
1960
1961#[cfg(test)]
1962mod test {
1963    use super::*;
1964    use num::complex::Complex32;
1965
1966    #[test]
1967    fn playground() {
1968        let a = arange((2.5, 3.2, 0.02));
1969        println!("{a:6.3?}");
1970        let a = arange(15.0);
1971        println!("{a:6.3?}");
1972        let a = arange((15.0, &DeviceCpu::default()));
1973        println!("{a:6.3?}");
1974        let a: Tensor<f64, _> = unsafe { empty(([15, 18].f(), &DeviceCpuSerial::default())) };
1975        println!("{a:6.3?}");
1976        let a = unsafe { a.empty_like() };
1977        println!("{a:6.3?}");
1978        let a = unsafe { empty_like((&a, TensorIterOrder::C)) };
1979        println!("{a:6.3?}");
1980        let a: Tensor<f64, _> = eye(3);
1981        println!("{a:6.3?}");
1982        let a = full(([2, 2].f(), 3.16));
1983        println!("{a:6.3?}");
1984        let a = full_like((&a, 2.71));
1985        println!("{a:6.3?}");
1986        let a = a.full_like(2.71);
1987        println!("{a:6.3?}");
1988        let a = linspace((3.2, 4.7, 12));
1989        println!("{a:6.3?}");
1990        let a = linspace((Complex32::new(1.8, 7.5), Complex32::new(-8.9, 1.6), 12));
1991        println!("{a:6.3?}");
1992        let a: Tensor<f64> = ones(vec![2, 2]);
1993        println!("{a:6.3?}");
1994        let a = a.ones_like();
1995        println!("{a:6.3?}");
1996        let a: Tensor<f64> = zeros([2, 2]);
1997        println!("{a:6.3?}");
1998        let a: Tensor<f64, _> = zeros(([2, 2], &DeviceCpuSerial::default()));
1999        println!("{a:6.3?}");
2000        let a = a.zeros_like();
2001        println!("{a:6.3?}");
2002    }
2003
2004    #[test]
2005    fn test_tril() {
2006        let a = arange((1, 10)).into_shape((3, 3));
2007        let b = a.view().tril();
2008        println!("{b:6.3?}");
2009        let b = triu((a, 1));
2010        println!("{b:6.3?}");
2011    }
2012}