Method HandleAllToAll

tensorflow/compiler/xla/service/hlo_verifier.cc:230–272

}

Status ShapeVerifier::HandleAllToAll(HloInstruction* hlo) {
  TF_RETURN_IF_ERROR(CheckReplicaGroups(hlo));

  auto* all_to_all = Cast<HloAllToAllInstruction>(hlo);
  TF_RET_CHECK(all_to_all != nullptr);
  if (all_to_all->split_dimension()) {
    if (hlo->replica_groups().empty()) {
      return InternalError(
          "An array all-to-all must have an explicit replica_groups config");
    }
  }

  // The size of each replica group must be the same (the split count of the
  // operaion). In case the default replica group is used (empty replica group,
  // must not be an array all-to-all, as checked above), infer from the number
  // of operands.
  const int64 split_count = hlo->replica_groups().empty()
                                ? hlo->operand_count()
                                : hlo->replica_groups()[0].replica_ids_size();
  for (const ReplicaGroup& g : hlo->replica_groups()) {
    if (g.replica_ids_size() != split_count) {
      return InternalError(
          "Replica group has size %d, but all replica groups in an all-to-all "
          "must have size N: %s",
          g.replica_ids_size(), hlo->ToString());
    }
  }

  if (all_to_all->split_dimension()) {
    TF_RET_CHECK(hlo->operand_count() == 1);
    return CheckShape(
        hlo, ShapeInference::InferAllToAllShape(
                 hlo->operand(0)->shape(), *all_to_all->split_dimension(),
                 *all_to_all->split_dimension(), split_count));
  } else {
    std::vector<const Shape*> operand_shapes;
    for (const HloInstruction* operand : hlo->operands()) {
      operand_shapes.push_back(&operand->shape());
    }
    return CheckShape(hlo,
                      ShapeInference::InferAllToAllTupleShape(operand_shapes));
  }
}

Status ShapeVerifier::HandlePartitionId(HloInstruction* hlo) {
  return CheckShape(hlo, ShapeUtil::MakeShape(U32, {}));