GetPartitionedOutputsAndIdentityOps(cluster);

  for (auto [_, ops] : partitioned_outputs) {
    if (!(ops.empty() || ops.size() == num_replicas)) {
      return (ops.begin())->emitOpError()
             << "expected zero or " << num_replicas
             << " 'TPUPartitionedOutput' op(s), instead got "
             << partitioned_outputs.size();
    }
  }

  // No need to replicate.

        base_dilations_vec, window_dilations_vec);
    if (!window) {
      op->emitOpError("failed to create window");
    }
    auto output_shape = InferWindowOutputShape(
        input_ty, window.value(),
        mlir::cast<ShapedType>(op.getInitValue().getType()).getElementType());

    if (!output_shape) {
      op->emitOpError("failed to infer output shape");
    }

        mlir::cast<FlatSymbolRefAttr>(sym_ref).getValue());

    if (!init_func_op)
      return session_initializer.emitOpError()
             << "the initializer function does not exist";

    if (!init_func_op.getFunctionType().getResults().empty())
      return session_initializer.emitOpError()
             << "the initializer function should have no output";

    auto exported_names = GetExportedNames(init_func_op);

      auto ret = func.front().getTerminator();
      auto block_arg =
          mlir::dyn_cast<BlockArgument>(ret->getOperand(result_idx));
      if (!block_arg) {
        return op->emitOpError("result #")
               << result_idx << " not tied to function argument for branch @"
               << func.getName();
      }
      if (!common_arg_index.has_value()) {

      }
      if (axis < 0 || (input_type.hasRank() && axis >= input_type.getRank())) {
        return transpose_op.emitOpError("perm must be in [-rank, rank)");
      }
      if (std::count(axes.begin(), axes.end(), axis) > 0) {
        return transpose_op.emitOpError("perm cannot have duplicated axis");
      }
      axes.push_back(axis);
    }

  if (failed(tensorflow::GetDevicesFromOp(moduleOp, &devices)))
    return moduleOp.emitOpError() << "No available devices.";
  llvm::ArrayRef<tensorflow::DeviceNameUtils::ParsedName> device_names =
      devices.device_names();
  auto status_or_system_devices = GetTPUSystemDevices(device_names);
  if (!status_or_system_devices.ok())
    return moduleOp.emitOpError()
           << "error in fetching TPU_SYSTEM devices: "

  if (!ranked_type)
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output types must be "
              "ranked tensors.";
  ArrayRef<int64_t> in_shape = ranked_type.getShape();
  if (in_shape.empty() || in_shape[0] < 0) {
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output shapes must "

  if (!num_cores_per_replica_attr)
    return cluster_func.emitOpError(
        CreateMissingAttributeMsg(tensorflow::kNumCoresPerReplicaAttr));

  int num_cores_per_replica = num_cores_per_replica_attr.getInt();

  auto topology_attr =
      cluster_func->getAttrOfType<StringAttr>(tensorflow::kTopologyAttr);
  if (!topology_attr)
    return cluster_func.emitOpError(
        CreateMissingAttributeMsg(tensorflow::kTopologyAttr));

    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||
        llvm::any_of(op->getResultTypes(), IsQuantType) ||
        IsTFUniformQuantizedOp(op) || IsMhloUniformQuantizedOp(*op)) {
      op->emitOpError("is illegal as it is a UQ op or contains uq/qint types");
      LOG(ERROR) << "Found illegal op containing uq/qint type: "
                 << op->getName().getStringRef().str();
      return WalkResult::interrupt();
    }

                                    Value* buffer) {
  auto max_count_op = max_size.getDefiningOp();
  if (!max_count_op) return op->emitOpError("unknown max element count");
  auto max_count_const_op = llvm::dyn_cast<TF::ConstOp>(max_count_op);
  if (!max_count_const_op) return op->emitOpError("unknown max element count");
  int64_t max_size_const =
      (*max_count_const_op.getValue().getValues<APInt>().begin())