Function ConsecutiveSegments

tensorflow/compiler/xla/shape_util.cc:1487–1495

Returns the indices of the first elements of all consecutive subarrays of the given array. For example: ConsecutiveSegments({m, m+1, m+2, n, k, k+1}) = {0, 3, 4}

// given array. For example:
// ConsecutiveSegments({m, m+1, m+2, n, k, k+1}) = {0, 3, 4}
static std::vector<size_t> ConsecutiveSegments(absl::Span<const int64> xs) {
  std::vector<size_t> is = {0};
  for (size_t i = 1; i < xs.size(); ++i) {
    if (1 != xs[i] - xs[i - 1]) {
      is.push_back(i);
    }
  }
  return is;
}

// Merges the sequences of dimensions of the given shape which start at the
// given indices `segs`.