SplitWithSizes

Splits the tensor into chunks chunks with specified sizes.

Abstract Signature:

SplitWithSizes(tensor: Tensor, split_sizes: List[int], dim: int = 0)

PyTorchJAX (Core)TensorFlow

PyTorch

API: torch.split_with_sizes

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: —

Strategy: Plugin (torch_split_with_sizes_converter)

TensorFlow

API: tf.split

Strategy: Direct Mapping

Official Docs ↗