burn::tensor

Function split

pub fn split<B, K>(
    tensor: <K as TensorKind<B>>::Primitive,
    split_size: usize,
    dim: usize,
) -> Vec<<K as TensorKind<B>>::Primitive>
where
    B: Backend,
    K: TensorKind<B> + BasicOps<B>,
Expand description

Splits the tensor along the given dimension into equally sized chunks (if possible) with size split_size. Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

§Arguments

  • tensor - The tensor.
  • split_size - The size of a single chunk.
  • dim - The dimension along which to split the tensor.

§Returns

A vector of tensors.

§Remarks

This (and the following) are fallback solutions that is used only when the backend doesn’t have the corresponding implementation. Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved by static dispatch. It is not designed for direct usage by users, and not recommended to import or use this function directly.