Function InsertBroadcastLabels

tensorflow/core/kernels/einsum_op.cc:178–187

Insert new (unnamed) broadcasting labels at the location of ellipsis.

// Insert new (unnamed) broadcasting labels at the location of ellipsis.
void InsertBroadcastLabels(int num_bcast_dims, int num_named_labels,
                           int ellipsis_axis, Labels* labels,
                           LabelCounts* label_counts) {
  labels->erase(labels->begin() + ellipsis_axis);
  labels->insert(labels->begin() + ellipsis_axis, num_bcast_dims, 0);
  std::iota(labels->begin() + ellipsis_axis,
            labels->begin() + ellipsis_axis + num_bcast_dims, num_named_labels);
  // Increment label counts. Since these are new labels, the count is set to 1.
  label_counts->resize(num_named_labels + num_bcast_dims, 1);
}

// Record and validate the label to dimension mapping. Must be a named
// (non-broadcasting) label as broadcasting labels don't have a fixed dimension.