Function SplitEasyCases

tensorflow/core/kernels/batch_kernels.cc:107–140

// applicable special case and wrote to the outputs. Otherwise acts as a no-op.
template <typename T>
Status SplitEasyCases(OpKernelContext* context, const Tensor& input,
                      const gtl::ArraySlice<int64>& sizes,
                      std::vector<Tensor>* outputs, bool* done) {
  *done = false;

  int64 total_size = 0;
  for (const int64 size : sizes) {
    total_size += size;
  }
  if (total_size > input.shape().dim_size(0)) {
    return errors::InvalidArgument(
        "Sum of split sizes must not exceed dim0-size of input tensor");
  }

  // Special case 0: trivial 1-way split.
  if (sizes.size() == 1 && sizes.at(0) == input.shape().dim_size(0)) {
    outputs->push_back(input);
    *done = true;
    return Status::OK();
  }

  // Special case 1: input is aligned.
  if (IsInnerDimsSizeAligned<T>(input.shape())) {
    int64 position = 0;
    for (const int64 size : sizes) {
      outputs->emplace_back(input.Slice(position, position + size));
      position += size;
    }
    *done = true;
    return Status::OK();
  }

  return Status::OK();
}

// Handles the general case, on CPU.
template <typename T>