| 3986 | # Note: does not work with TensorMetas because of data-dependent control-flow |
| 3987 | # CompositeImplicitAutograd - don't register decomp |
| 3988 | def tensor_split( |
| 3989 | a: TensorLikeType, |
| 3990 | indices_or_sections: Union[Tensor, DimsType], |
| 3991 | dim: int = 0, |
| 3992 | ) -> Tuple[TensorLikeType, ...]: |
| 3993 | _dim = utils.canonicalize_dim(a.ndim, dim) |
| 3994 | if a.ndim == 0: |
| 3995 | msg = "tensor_split: received a rank zero tensor, but expected a tensor of rank one or greater!" |
| 3996 | raise ValueError(msg) |
| 3997 | |
| 3998 | # If indices_or_sections is a tensor, it must be a CPU Long tensor |
| 3999 | if isinstance(indices_or_sections, TensorLike): |
| 4000 | if not indices_or_sections.device.type == "cpu": |
| 4001 | msg = "tensor_split: if indices_or_sections is a tensor it must be on the CPU, but received one on {}".format( |
| 4002 | indices_or_sections.device |
| 4003 | ) |
| 4004 | raise ValueError(msg) |
| 4005 | if indices_or_sections.dtype != torch.long: |
| 4006 | msg = "tensor_split: if indices_or_sections is a tensor it must have long dtype, " |
| 4007 | f" but received one with dtype {indices_or_sections.dtype}" |
| 4008 | raise ValueError(msg) |
| 4009 | |
| 4010 | # Case 0 -- indices_or_sections is an integer or a scalar tensor n and a is split along dim into n parts of equal-ish length |
| 4011 | if isinstance(indices_or_sections, IntLike) or ( |
| 4012 | isinstance(indices_or_sections, TensorLike) and indices_or_sections.ndim == 0 |
| 4013 | ): |
| 4014 | sections: int = ( |
| 4015 | indices_or_sections # type: ignore[assignment] |
| 4016 | if isinstance(indices_or_sections, Number) |
| 4017 | else indices_or_sections.item() |
| 4018 | ) |
| 4019 | |
| 4020 | if sections <= 0: |
| 4021 | msg = f"tensor_split: number of sections must be greater than 0, but was {sections}" |
| 4022 | raise ValueError(msg) |
| 4023 | |
| 4024 | splits = [] |
| 4025 | dim_size = a.shape[_dim] |
| 4026 | min_split_size = math.floor(dim_size / sections) |
| 4027 | num_splits_one_extra = dim_size % sections |
| 4028 | start_idx = 0 |
| 4029 | for split_idx in range(sections): |
| 4030 | split_size = ( |
| 4031 | min_split_size + 1 |
| 4032 | if (split_idx < num_splits_one_extra) |
| 4033 | else min_split_size |
| 4034 | ) |
| 4035 | s = prims.slice_in_dim(a, start_idx, start_idx + split_size, axis=_dim) |
| 4036 | splits.append(s) |
| 4037 | start_idx = start_idx + split_size |
| 4038 | |
| 4039 | return tuple(splits) |
| 4040 | # Case 1 -- indices_or_sections is a sequence of integers or a 1D tensor describing the splits |
| 4041 | else: |
| 4042 | indices = indices_or_sections |
| 4043 | if isinstance(indices_or_sections, TensorLike): |
| 4044 | if indices_or_sections.ndim != 1: |
| 4045 | msg = "tensor_split: non-scalar indices_or_sections tensors must have only one dimension, " |