Function burn::tensor::cartesian_grid

pub fn cartesian_grid<B, S, const D: usize, const D2: usize>(
    shape: S,
    device: &<B as Backend>::Device,
) -> <B as Backend>::IntTensorPrimitive<D2>
where
    B: Backend,
    S: Into<Shape<D>>,
Expand description

Generates a cartesian grid for the given tensor shape on the specified device. The generated tensor is of dimension D2 = D + 1, where each element at dimension D contains the cartesian grid coordinates for that element.

§Arguments

  • shape - The shape specifying the dimensions of the tensor.
  • device - The device to create the tensor on.

§Panics

Panics if D2 is not equal to D+1.

§Examples

   use burn_tensor::Int;
   use burn_tensor::{backend::Backend, Shape, Tensor};
   fn example<B: Backend>() {
       let device = Default::default();
       let result: Tensor<B, 3, _> = Tensor::<B, 2, Int>::cartesian_grid([2, 3], &device);
       println!("{}", result);
   }