Function tensor_split

paddle/phi/api/include/compat/ATen/ops/tensor_split.h:25–64

namespace at {

inline std::vector<at::Tensor> tensor_split(const at::Tensor& self,
                                            int64_t sections,
                                            int64_t dim = 0) {
  // Follow PyTorch's tensor_split_sections_symint implementation
  PD_CHECK(self._PD_GetInner().dims().size() > 0,
           "tensor_split expected at least a 1-dimensional tensor, but got a "
           "tensor with ",
           self._PD_GetInner().dims().size(),
           " dims");

  PD_CHECK(
      sections > 0, "number of sections must be larger than 0, got ", sections);

  int64_t dim_size = self._PD_GetInner().dims()[dim];

  // Calculate split sizes: first (dim_size % sections) chunks get size
  // (dim_size / sections + 1), remaining chunks get size (dim_size / sections)
  auto min_split_size = dim_size / sections;
  auto num_splits_one_extra = dim_size % sections;

  std::vector<int64_t> split_sizes;
  split_sizes.reserve(sections);

  for (int64_t split_idx = 0; split_idx < sections; ++split_idx) {
    auto split_size = (split_idx < num_splits_one_extra) ? (min_split_size + 1)
                                                         : min_split_size;
    split_sizes.push_back(split_size);
  }

  // Use split with calculated sizes
  auto outputs =
      paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);

  std::vector<at::Tensor> at_tensors;
  at_tensors.reserve(outputs.size());
  for (const auto& paddle_tensor : outputs) {
    at_tensors.emplace_back(paddle_tensor);
  }
  return at_tensors;
}

inline std::vector<at::Tensor> tensor_split_symint(const at::Tensor& self,
                                                   c10::SymInt sections,