rstsr_core/tensor/
adv_indexing.rs

1//! Advanced indexing related tensor manipulations.
2//!
3//! Currently, full support of advanced indexing is not available. However, it
4//! is still possible to index one axis by list.
5
6use crate::prelude_dev::*;
7
8/* #region index_select */
9
10pub fn index_select_f<R, T, B, D, I>(tensor: &TensorAny<R, T, B, D>, axis: isize, indices: I) -> Result<Tensor<T, B, D>>
11where
12    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
13    D: DimAPI + DimSmallerOneAPI,
14    D::SmallerOne: DimAPI,
15    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
16    I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
17{
18    // TODO: output layout control (TensorIterOrder::K or default layout)
19    let device = tensor.device().clone();
20    let tensor_layout = tensor.layout();
21    let ndim = tensor_layout.ndim();
22    // check axis and index
23    let axis = rstsr_check_axis!(axis, ndim)?;
24    let nshape: usize = tensor_layout.shape()[axis];
25    let indices = indices.try_into().map_err(Into::into)?;
26    let indices = indices
27        .as_ref()
28        .iter()
29        .map(|&i| -> Result<usize> {
30            let i = if i < 0 { nshape as isize + i } else { i };
31            rstsr_pattern!(
32                i,
33                0..nshape as isize,
34                IndexError,
35                "Invalid index that exceeds shape length at axis {}.",
36                axis
37            )?;
38            Ok(i as usize)
39        })
40        .collect::<Result<Vec<usize>>>()?;
41    let mut out_shape = tensor_layout.shape().as_ref().to_vec();
42    out_shape[axis] = indices.len();
43    let out_layout = out_shape.new_contig(None, device.default_order()).into_dim()?;
44    let mut out_storage = device.uninit_impl(out_layout.size())?;
45    device.index_select(out_storage.raw_mut(), &out_layout, tensor.storage().raw(), tensor_layout, axis, &indices)?;
46    let out_storage = unsafe { B::assume_init_impl(out_storage)? };
47    TensorBase::new_f(out_storage, out_layout)
48}
49
50/// Returns a new tensor, which indexes the input tensor along dimension `axis`
51/// using the entries in `indices`.
52///
53/// # See also
54///
55/// This function should be similar to PyTorch's [`torch.index_select`](https://docs.pytorch.org/docs/stable/generated/torch.index_select.html).
56pub fn index_select<R, T, B, D, I>(tensor: &TensorAny<R, T, B, D>, axis: isize, indices: I) -> Tensor<T, B, D>
57where
58    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
59    D: DimAPI + DimSmallerOneAPI,
60    D::SmallerOne: DimAPI,
61    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
62    I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
63{
64    index_select_f(tensor, axis, indices).rstsr_unwrap()
65}
66
67pub fn take_f<R, T, B, D, I>(tensor: &TensorAny<R, T, B, D>, indices: I, axis: isize) -> Result<Tensor<T, B, D>>
68where
69    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
70    D: DimAPI + DimSmallerOneAPI,
71    D::SmallerOne: DimAPI,
72    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
73    I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
74{
75    index_select_f(tensor, axis, indices)
76}
77
78/// Take elements from an array along an axis.
79///
80/// # See also
81///
82/// [Python Array API standard: take](https://data-apis.org/array-api/latest/API_specification/generated/array_api.take.html#array_api.take)
83pub fn take<R, T, B, D, I>(tensor: &TensorAny<R, T, B, D>, indices: I, axis: isize) -> Tensor<T, B, D>
84where
85    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
86    D: DimAPI + DimSmallerOneAPI,
87    D::SmallerOne: DimAPI,
88    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
89    I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
90{
91    index_select(tensor, axis, indices)
92}
93
94impl<R, T, B, D> TensorAny<R, T, B, D>
95where
96    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
97    D: DimAPI + DimSmallerOneAPI,
98    D::SmallerOne: DimAPI,
99    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
100{
101    pub fn index_select_f<I>(&self, axis: isize, indices: I) -> Result<Tensor<T, B, D>>
102    where
103        I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
104    {
105        index_select_f(self, axis, indices)
106    }
107
108    /// Returns a new tensor, which indexes the input tensor along dimension
109    /// `axis` using the entries in `indices`.
110    ///
111    /// # See also
112    ///
113    /// This function should be similar to PyTorch's [`torch.index_select`](https://docs.pytorch.org/docs/stable/generated/torch.index_select.html).
114    pub fn index_select<I>(&self, axis: isize, indices: I) -> Tensor<T, B, D>
115    where
116        I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
117    {
118        index_select(self, axis, indices)
119    }
120
121    pub fn take_f<I>(&self, indices: I, axis: isize) -> Result<Tensor<T, B, D>>
122    where
123        I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
124    {
125        take_f(self, indices, axis)
126    }
127
128    /// Take elements from an array along an axis.
129    ///
130    /// # See also
131    ///
132    /// [Python Array API standard: take](https://data-apis.org/array-api/latest/API_specification/generated/array_api.take.html#array_api.take)
133    pub fn take<I>(&self, indices: I, axis: isize) -> Tensor<T, B, D>
134    where
135        I: TryInto<AxesIndex<isize>, Error: Into<Error>>,
136    {
137        take(self, indices, axis)
138    }
139}
140
141/* #endregion */
142
143/* #region bool_select */
144
145pub fn bool_select_f<R, T, B, D, I>(tensor: &TensorAny<R, T, B, D>, axis: isize, mask: I) -> Result<Tensor<T, B, D>>
146where
147    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
148    D: DimAPI + DimSmallerOneAPI,
149    D::SmallerOne: DimAPI,
150    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
151    I: TryInto<AxesIndex<bool>, Error: Into<Error>>,
152{
153    // transform bool to index
154    let indices = mask
155        .try_into()
156        .map_err(Into::into)?
157        .as_ref()
158        .iter()
159        .enumerate()
160        .filter_map(|(i, &m)| m.then_some(i))
161        .collect::<Vec<usize>>();
162    index_select_f(tensor, axis, indices)
163}
164
165/// Returns a new tensor, which indexes the input tensor along dimension `axis`
166/// using the boolean entries in `mask`.
167pub fn bool_select<R, T, B, D, I>(tensor: &TensorAny<R, T, B, D>, axis: isize, mask: I) -> Tensor<T, B, D>
168where
169    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
170    D: DimAPI + DimSmallerOneAPI,
171    D::SmallerOne: DimAPI,
172    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
173    I: TryInto<AxesIndex<bool>, Error: Into<Error>>,
174{
175    bool_select_f(tensor, axis, mask).rstsr_unwrap()
176}
177
178impl<R, T, B, D> TensorAny<R, T, B, D>
179where
180    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
181    D: DimAPI + DimSmallerOneAPI,
182    D::SmallerOne: DimAPI,
183    B: DeviceAPI<T> + DeviceIndexSelectAPI<T, D> + DeviceCreationAnyAPI<T>,
184{
185    pub fn bool_select_f<I>(&self, axis: isize, indices: I) -> Result<Tensor<T, B, D>>
186    where
187        I: TryInto<AxesIndex<bool>, Error: Into<Error>>,
188    {
189        bool_select_f(self, axis, indices)
190    }
191
192    /// Returns a new tensor, which indexes the input tensor along dimension
193    /// `axis` using the boolean entries in `mask`.
194    pub fn bool_select<I>(&self, axis: isize, indices: I) -> Tensor<T, B, D>
195    where
196        I: TryInto<AxesIndex<bool>, Error: Into<Error>>,
197    {
198        bool_select(self, axis, indices)
199    }
200}
201
202/* #endregion */
203
204#[cfg(test)]
205mod test {
206    use super::*;
207
208    #[test]
209    fn test_index_select() {
210        #[cfg(not(feature = "col_major"))]
211        {
212            let device = DeviceCpuSerial::default();
213            let a = linspace((1.0, 24.0, 24, &device)).into_shape((2, 3, 4));
214            let b = a.index_select(0, [0, 0, 1, -1]);
215            assert!(fingerprint(&b) - -31.94175930917264 < 1e-8);
216            let b = a.index_select(1, [0, 0, 1, -1]);
217            assert!(fingerprint(&b) - 3.5719025258942088 < 1e-8);
218            let b = a.index_select(2, [0, 0, 1, -1]);
219            assert!(fingerprint(&b) - -25.648600916145096 < 1e-8);
220        }
221        #[cfg(feature = "col_major")]
222        {
223            let device = DeviceCpuSerial::default();
224            let a = linspace((1.0, 24.0, 24, &device)).into_shape((4, 3, 2));
225            let b = a.index_select(2, [0, 0, 1, -1]);
226            assert!(fingerprint(&b) - -31.94175930917264 < 1e-8);
227            let b = a.index_select(1, [0, 0, 1, -1]);
228            assert!(fingerprint(&b) - 3.5719025258942088 < 1e-8);
229            let b = a.index_select(0, [0, 0, 1, -1]);
230            assert!(fingerprint(&b) - -25.648600916145096 < 1e-8);
231        }
232
233        // 1-dim select with empty index
234        let device = DeviceCpuSerial::default();
235        let a = linspace((1.0, 4.0, 4, &device));
236        let mask: Vec<usize> = vec![];
237        let b = a.index_select(0, &mask);
238        assert_eq!(b.raw(), &[]);
239    }
240
241    #[test]
242    fn test_index_select_default_device() {
243        #[cfg(not(feature = "col_major"))]
244        {
245            let device = DeviceCpu::default();
246            let a = linspace((1.0, 2.0, 256 * 256 * 256, &device)).into_shape((256, 256, 256));
247            let sel = [1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233];
248            let b = a.index_select(0, sel);
249            assert!(fingerprint(&b) - 0.9357016252766746 < 1e-10);
250            let b = a.index_select(1, sel);
251            assert!(fingerprint(&b) - 1.012193909979973 < 1e-10);
252            let b = a.index_select(2, sel);
253            assert!(fingerprint(&b) - 1.010735112247236 < 1e-10);
254        }
255        #[cfg(feature = "col_major")]
256        {
257            let device = DeviceCpu::default();
258            let a = linspace((1.0, 2.0, 256 * 256 * 256, &device)).into_shape((256, 256, 256));
259            let sel = [1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233];
260            let b = a.index_select(2, sel);
261            assert!(fingerprint(&b) - 0.9357016252766746 < 1e-10);
262            let b = a.index_select(1, sel);
263            assert!(fingerprint(&b) - 1.012193909979973 < 1e-10);
264            let b = a.index_select(0, sel);
265            assert!(fingerprint(&b) - 1.010735112247236 < 1e-10);
266        }
267    }
268
269    #[test]
270    fn test_bool_select_workable() {
271        let a = arange(24).into_shape((2, 3, 4));
272        let b = a.bool_select(-2, [true, false, true]);
273        println!("{b:?}");
274    }
275}
rstsr_core/tensor/adv_indexing.rs

rstsr_core/tensor/
adv_indexing.rs
