burn_tensor/tensor/api/
cartesian_grid.rs

1use crate::{Device, Int, Shape, Tensor};
2use alloc::vec::Vec;
3
4/// Generates a cartesian grid for the given tensor shape on the specified device.
5/// The generated tensor is of dimension `D2 = D + 1`, where each element at dimension D contains the cartesian grid coordinates for that element.
6///
7/// # Arguments
8///
9/// * `shape` - The shape specifying the dimensions of the tensor.
10/// * `device` - The device to create the tensor on.
11///
12/// # Panics
13///
14/// Panics if `D2` is not equal to `D+1`.
15///
16/// # Examples
17///
18/// ```rust
19///    use burn_tensor::Int;
20///    use burn_tensor::{Shape, Tensor};
21/// let device = Default::default();
22/// let result: Tensor<3, _> = Tensor::<2, Int>::cartesian_grid([2, 3], &device);
23/// println!("{}", result);
24/// ```
25pub fn cartesian_grid<S: Into<Shape>, const D: usize, const D2: usize>(
26    shape: S,
27    device: &Device,
28) -> Tensor<D2, Int> {
29    if D2 != D + 1 {
30        panic!("D2 must equal D + 1 for Tensor::cartesian_grid")
31    }
32
33    let dims = shape.into();
34    let mut indices: Vec<Tensor<D, Int>> = Vec::new();
35
36    for dim in 0..D {
37        let dim_range: Tensor<1, Int> = Tensor::arange(0..dims[dim] as i64, device);
38
39        let mut shape = [1; D];
40        shape[dim] = dims[dim];
41        let mut dim_range = dim_range.reshape(shape);
42
43        for (i, &item) in dims.iter().enumerate() {
44            if i == dim {
45                continue;
46            }
47            dim_range = dim_range.repeat_dim(i, item);
48        }
49
50        indices.push(dim_range);
51    }
52
53    Tensor::stack::<D2>(indices, D)
54}
burn_tensor/tensor/api/cartesian_grid.rs

burn_tensor/tensor/api/
cartesian_grid.rs
