Function ComputeStrides

tensorflow/compiler/xla/service/hlo_evaluator.cc:950–970

Helper to compute strides for creating linear indices into multidimensional data from the dimension lengths and the layout. Returns a new vector of size lengths.size() + 1. The last element of the returned vector at index [lengths.size()] contains the product of all dimension lengths.

// lengths.size() + 1. The last element of the returned vector at index
// [lengths.size()] contains the product of all dimension lengths.
std::vector<int64> ComputeStrides(const absl::Span<const int64> lengths,
                                  const Layout& layout) {
  const int64 num_dimensions = lengths.size();

  // Make sure that the layout length matches the number of dimensions.
  CHECK_EQ(num_dimensions, layout.minor_to_major_size());

  // Calculate strides using layout-specified ordering of the dimensions and
  // place the stride for axis 0 at index 0, for axis 1 at index 1, etc.
  std::vector<int64> strides(num_dimensions + 1);
  int64 stride = 1;
  for (int64 i = 0; i < num_dimensions; i++) {
    // Reverse the ordering of the dimensions in the layout.
    const int64 index = (num_dimensions - 1) - layout.minor_to_major(i);
    strides[index] = stride;
    stride *= lengths[index];
  }
  strides[num_dimensions] = stride;

  return strides;
}

// Compute strides as above using the default layout.
std::vector<int64> ComputeStrides(const absl::Span<const int64> lengths) {