ttnn.split

ttnn.split() → List[ttnn.Tensor]

Splits input_tensor into chunks along dimension dim and returns them as a list of tensors.

The behavior depends on the type of split_size:

  • If split_size is an int, the tensor is split into contiguous chunks of split_size elements along dim. When input_tensor.shape[dim] is not an exact multiple of split_size, the final chunk holds the remainder and is smaller. The number of outputs is ceil(input_tensor.shape[dim] / split_size).

  • If split_size is a list[int], the tensor is split into len(split_size) contiguous chunks whose sizes along dim are the list entries, in order. The entries must sum exactly to input_tensor.shape[dim].

Constraints:

  • Every chunk size must be greater than 0 (both a zero entry in a split_size list and a zero-size split dimension raise an error; zero-volume tensors are not supported by the device kernels).

  • For a split_size list, the entries must sum exactly to input_tensor.shape[dim]; an error is raised for both under- and over-covering lists.

Example:

# int split_size: 6 along dim=1 -> chunks of size 2
a, b, c = ttnn.split(input_tensor, 2, dim=1)  # input_tensor.shape[1] == 6

# list split_size: explicit per-chunk sizes summing to shape[1]
a, b = ttnn.split(input_tensor, [2, 4], dim=1)  # input_tensor.shape[1] == 6

:param * input_tensor: Input Tensor. :param * split_size: Size of a single chunk, or a list of per-chunk sizes. :type * split_size: Union[int, list[int]] :param * dim: Dimension along which to split. Negative indexing is supported. Defaults to 0. :type * dim: int

:keyword * memory_config: Memory Config of the output tensors.

Returns:

List[ttnn.Tensor]The list of output tensors.