Skip to main content

Mojo function

split

split[n: Int](input: Symbol, sizes: StaticIntTuple[n], axis: Int = 0) -> List[Symbol, 0]

Splits a symbolic tensor into specified bucket sizes along the axis.

Parameters:

  • n (Int): The number of symbolic tensors to split into.

Args:

  • input (Symbol): The symbolic tensor to split.
  • sizes (StaticIntTuple[n]): The list of sizes for each split.
  • axis (Int): The axis to split along.

Returns:

n symbolic tensor values. The ith result corresponds to the sizes[i]. Each tensor will have the same rank as the input, and will be the result of slicing the input of sizes[i] elements along axis, starting at the offset of the cumulative sum of the previous sizes.