Function ProcessDimensions

tensorflow/core/kernels/einsum_op.cc:208–290

Validate input dimensions and populate unnamed labels and their label counts.

// Validate input dimensions and populate unnamed labels and their label counts.
Status ProcessDimensions(const OpInputList& inputs,
                         const gtl::InlinedVector<bool, 2>& input_has_ellipsis,
                         const bool output_has_ellipsis,
                         OperandLabels* input_labels, Labels* output_labels,
                         std::vector<DimensionType>* label_types,
                         OperandLabelCounts* input_label_counts,
                         LabelCounts* output_label_counts,
                         LabelToDimSizes* label_to_dim_sizes) {
  if (inputs.size() != input_labels->size()) {
    return errors::InvalidArgument("Expected ", input_labels->size(),
                                   " inputs but got: ", inputs.size());
  }
  const int num_inputs = inputs.size();

  // We infer the number of broadcasting dimensions by taking the maximum rank
  // among the broadcasting subshapes of the input.
  int max_bcast_dims = 0;
  const int num_named_labels = label_types->size();
  label_to_dim_sizes->resize(num_named_labels);
  for (int i = 0; i < num_inputs; ++i) {
    Labels* labels = &(*input_labels)[i];

    if (!input_has_ellipsis[i]) {
      if (inputs[i].dims() != labels->size()) {
        return errors::InvalidArgument("Expected input ", i, " to have rank ",
                                       labels->size(),
                                       " but got: ", inputs[i].dims());
      }
      for (int label_idx = 0; label_idx < labels->size(); ++label_idx) {
        const int label = (*labels)[label_idx];
        TF_RETURN_IF_ERROR(RecordLabelToDimension(label, label_idx, inputs[i],
                                                  label_to_dim_sizes));
      }
      continue;
    }

    // Input has an ellipsis.
    if (inputs[i].dims() + 1 < labels->size()) {
      return errors::InvalidArgument(
          "Expected input ", i, " to have rank at least ", labels->size() - 1,
          " but got: ", inputs[i].dims());
    }
    int ellipsis_axis = -1;
    const int num_bcast_dims = inputs[i].dims() - labels->size() + 1;
    for (int label_idx = 0; label_idx < labels->size(); ++label_idx) {
      const int label = (*labels)[label_idx];
      if (label == kEllipsisLabel) {
        ellipsis_axis = label_idx;
        continue;
      }
      // Current label is not an ellipsis.
      const int axis =
          label_idx + (ellipsis_axis == -1 ? 0 : num_bcast_dims - 1);
      TF_RETURN_IF_ERROR(
          RecordLabelToDimension(label, axis, inputs[i], label_to_dim_sizes));
    }
    // Found an ellipsis. Replace 'kEllipsisLabel' with broadcasting dimensions.
    if (ellipsis_axis != -1) {