Method InferDotOpShape

tensorflow/compiler/xla/service/shape_inference.cc:627–718

static */

}  // namespace

/* static */ StatusOr<Shape> ShapeInference::InferDotOpShape(
    const Shape& lhs, const Shape& rhs,
    const DotDimensionNumbers& dimension_numbers) {
  TF_RETURN_IF_ERROR(ExpectArray(lhs, "lhs of dot"));
  TF_RETURN_IF_ERROR(ExpectArray(rhs, "rhs of dot"));

  auto fail = [lhs, rhs](const string& addendum) -> Status {
    string message =
        StrFormat("Cannot infer shape for dot operation: %s <dot> %s.",
                  ShapeUtil::HumanString(lhs), ShapeUtil::HumanString(rhs));
    if (!addendum.empty()) {
      message += " " + addendum;
    }
    return InvalidArgument("%s", message);
  };

  // Check if both element types are the same.
  if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) {
    return fail("Element types do not match.");
  }

  // Validate basic properties of dot dimension numbers.
  TF_RETURN_IF_ERROR(ValidateDotDimensionNumbers(lhs, rhs, dimension_numbers));

  // Check that number of contracting dimensions match.
  if (dimension_numbers.lhs_contracting_dimensions_size() !=
      dimension_numbers.rhs_contracting_dimensions_size()) {
    return fail(
        "Must specify the same number of contracting dimensions for lhs and "
        "rhs.");
  }
  // Check that contracting dimension sizes match.
  for (int64 i = 0; i < dimension_numbers.lhs_contracting_dimensions_size();
       ++i) {
    const int64 lhs_contracting_dimension =
        dimension_numbers.lhs_contracting_dimensions(i);
    const int64 rhs_contracting_dimension =
        dimension_numbers.rhs_contracting_dimensions(i);
    if (lhs.dimensions(lhs_contracting_dimension) !=
        rhs.dimensions(rhs_contracting_dimension)) {
      return fail("Contracting dimension sizes do not match.");
    }
  }

  // Check that number of batch dimensions match.
  if (dimension_numbers.lhs_batch_dimensions_size() !=
      dimension_numbers.rhs_batch_dimensions_size()) {
    return fail("Must the same number of batch dimensions for lhs and rhs.");
  }

  // Check that batch dimension numbers and sizes match.
  for (int64 i = 0; i < dimension_numbers.lhs_batch_dimensions_size(); ++i) {
    if (lhs.dimensions(dimension_numbers.lhs_batch_dimensions(i)) !=
        rhs.dimensions(dimension_numbers.rhs_batch_dimensions(i))) {
      return fail("Batch dimension sizes must match for lhs/rhs.");
    }
  }