burn_tensor/tensor/api/
take.rs

1use crate::{
2    AsIndex, Int, Tensor, check, check::TensorCheck, check::unwrap_dim_index, kind::Basic,
3};
4use alloc::vec::Vec;
5
6impl<const D: usize, K> Tensor<D, K>
7where
8    K: Basic,
9{
10    /// Takes elements from the tensor along the given dimension using indices of any dimensionality.
11    ///
12    /// This behaves like numpy's take function. When indices is multi-dimensional,
13    /// the output shape will be: input.shape\[:dim\] + indices.shape + input.shape\[dim+1:\]
14    ///
15    /// # Arguments
16    ///
17    /// * `dim` - The dimension along which to select elements. Supports negative indexing.
18    /// * `indices` - The indices of elements to select. Can be any dimensionality.
19    ///   Must be valid indices in the range [0, dim_size).
20    ///
21    /// # Example
22    ///
23    /// ```rust
24    /// use burn_tensor::{Tensor, Int};
25    ///
26    /// let device = Default::default();
27    ///
28    /// // Example with 1D indices
29    /// let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
30    /// let indices = Tensor::<1, Int>::from_data([2, 0, 1], &device);
31    /// let result: Tensor<2> = tensor.clone().take::<1, 2>(-1, indices);  // -1 refers to last dimension
32    /// println!("{result}");
33    /// // [[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]]
34    ///
35    /// // Example with 2D indices - output will have +1 dimension (2D -> 3D)
36    /// let indices_2d = Tensor::<2, Int>::from_data([[0, 2], [1, 0]], &device);
37    /// let result: Tensor<3> = tensor.take::<2, 3>(1, indices_2d);
38    /// println!("{result}");
39    /// // [[[1.0, 3.0], [2.0, 1.0]], [[4.0, 6.0], [5.0, 4.0]]]
40    /// ```
41    pub fn take<const DI: usize, const DO: usize>(
42        self,
43        dim: impl AsIndex,
44        indices: Tensor<DI, Int>,
45    ) -> Tensor<DO, K> {
46        let dim = unwrap_dim_index(dim.try_dim_index(D), "Take");
47        check!(TensorCheck::take::<D, DI, DO>());
48
49        // Store the indices shape for reshaping later
50        let indices_shape = indices.shape();
51        let indices_dims = indices_shape.clone();
52
53        // Flatten indices to 1D for processing
54        let indices_flat = indices.reshape([indices_shape.num_elements()]);
55
56        // Perform the selection with the flattened indices
57        let selected = self.select(dim, indices_flat);
58
59        // Build the output shape
60        // Output shape = input.shape[:dim] + indices.shape + input.shape[dim+1:]
61        let selected_shape = selected.shape();
62        let mut new_shape = Vec::with_capacity(DO);
63
64        // Add dimensions before the selected dimension
65        for i in 0..dim {
66            new_shape.push(selected_shape[i]);
67        }
68
69        // Add all indices dimensions
70        for &idx_dim in indices_dims.iter() {
71            new_shape.push(idx_dim);
72        }
73
74        // Add dimensions after the selected dimension
75        for i in (dim + 1)..D {
76            new_shape.push(selected_shape[i]);
77        }
78
79        // Verify we have the correct number of dimensions
80        assert_eq!(
81            new_shape.len(),
82            DO,
83            "Internal error: shape calculation resulted in {} dims but expected {}",
84            new_shape.len(),
85            DO
86        );
87
88        // Convert to fixed-size array for reshape
89        let mut shape_array = [0; DO];
90        for (i, &s) in new_shape.iter().enumerate() {
91            shape_array[i] = s;
92        }
93
94        selected.reshape(shape_array)
95    }
96}
burn_tensor/tensor/api/take.rs

burn_tensor/tensor/api/
take.rs
