Method InferReshapeShape

tensorflow/compiler/xla/service/shape_inference.cc:2840–2989

static */

}

/* static */ StatusOr<Shape> ShapeInference::InferReshapeShape(
    const Shape& operand, absl::Span<const int64> dimensions,
    absl::Span<const int64> new_sizes, int64 inferred_dimension) {
  TF_RETURN_IF_ERROR(ExpectArray(operand, "reshape"));

  Shape inferred_shape =
      ShapeUtil::MakeShape(operand.element_type(), new_sizes);
  VLOG(3) << "Reshape inferred shape: "
          << ShapeUtil::HumanString(inferred_shape);

  if (ShapeUtil::ElementsIn(operand) != ShapeUtil::ElementsIn(inferred_shape)) {
    return InvalidArgument(
        "Reshape operation has mismatched element counts: from=%d (%s) "
        "to=%d (%s).",
        ShapeUtil::ElementsIn(operand), ShapeUtil::HumanString(operand),
        ShapeUtil::ElementsIn(inferred_shape),
        ShapeUtil::HumanString(inferred_shape));
  }

  std::vector<int64> indices(operand.rank());
  std::iota(indices.begin(), indices.end(), 0);
  if (dimensions.size() != operand.rank() ||
      !std::is_permutation(dimensions.begin(), dimensions.end(),
                           indices.begin())) {
    return InvalidArgument(
        "Reshape dimensions [%s] are not a permutation of the operand "
        "dimensions (operand shape is %s).",
        StrJoin(dimensions, ","), ShapeUtil::HumanString(operand));
  }

  // Propagate dynamic dimension.
  auto common_factors = CommonFactors(operand.dimensions(), new_sizes);
  for (int64 input_dim = 0; input_dim < operand.rank(); ++input_dim) {
    if (!operand.is_dynamic_dimension(input_dim)) {
      continue;
    }

    string reshape_debug_str = absl::StrFormat(
        "output: %s, input: %s, input_dim: "
        "%lld",
        ShapeUtil::HumanString(inferred_shape), ShapeUtil::HumanString(operand),
        input_dim);

    int64 input_dim_start = -1;
    int64 input_dim_end = -1;
    int64 output_dim_start = -1;
    int64 output_dim_end = -1;
    // Find common_factors that the input_dim belongs to.
    for (int64 i = 0; i < common_factors.size() - 1; ++i) {
      auto start = common_factors[i];
      auto end = common_factors[i + 1];
      if (input_dim >= start.first && input_dim < end.first) {
        input_dim_start = start.first;
        input_dim_end = end.first;
        output_dim_start = start.second;
        output_dim_end = end.second;
        break;
      }