Method FindTranspose021

tensorflow/compiler/xla/shape_util.cc:1513–1549

static*/

}

/*static*/ absl::optional<std::vector<int64>> ShapeUtil::FindTranspose021(
    const Shape& a, const Shape& b) {
  if (!CompatibleIgnoringElementType(a, b)) {
    return absl::nullopt;
  }

  std::vector<int64> permutation(a.dimensions().size());
  absl::Span<const int64> minor_to_major_a = LayoutUtil::MinorToMajor(a);
  std::vector<int64> major_to_minor_a(minor_to_major_a.rbegin(),
                                      minor_to_major_a.rend());
  absl::Span<const int64> minor_to_major_b = LayoutUtil::MinorToMajor(b);
  std::vector<int64> major_to_minor_b(minor_to_major_b.rbegin(),
                                      minor_to_major_b.rend());
  for (size_t i = 0; i < permutation.size(); ++i) {
    permutation[i] = PositionInContainer(major_to_minor_b, major_to_minor_a[i]);
  }

  std::vector<size_t> segments = ConsecutiveSegments(permutation);
  if ((3 == segments.size() && 0 == permutation[0]) || 2 == segments.size()) {
    Shape descending_layout_shape =
        ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(a);
    Shape normalized_shape = MergeDimensions(segments, descending_layout_shape);
    absl::Span<const int64> normalized_dims =
        AsInt64Slice(normalized_shape.dimensions());
    std::vector<int64> dims_021;
    if (2 == segments.size()) {
      // The logical component-0 is of size one.
      dims_021 = {1, normalized_dims[1], normalized_dims[0]};
    } else {
      dims_021 = {normalized_dims[0], normalized_dims[2], normalized_dims[1]};
    }

    return dims_021;
  }

  return absl::nullopt;
}

}  // namespace xla