bool IsOpReplicateInvariant(Region* replicate_region, Operation* op) {
  auto ancestor_of_replicate = [&](Region* region) {
    return region && region->isProperAncestor(replicate_region);
  };

  for (Value operand : op->getOperands())
    if (!ancestor_of_replicate(operand.getParentRegion())) return false;

  // _TPUDeviceOrdinalPlaceholder implicitly depends on the replica.
  if (llvm::isa<TF::_TPUDeviceOrdinalPlaceholderOp>(op)) return false;

    cost_map_[op] = kDefaultCheapCost;
    return;
  }

  // For other ops, use the sum of input sizes as its cost.
  int64_t cost = kDefaultCheapCost;
  for (auto operand : op->getOperands()) {
    auto type = mlir::cast<mlir::TensorType>(operand.getType());
    if (type.hasRank()) {
      cost += GetRankedTensorSize(type);
    } else {

                                    OpBuilder* builder) {
  builder->create<func::ReturnOp>(cluster_return_op.getLoc(),
                                  cluster_return_op.getOperands());
  cluster_return_op.erase();
}

// Builds a function that outlines region attached to cluster_op or launch_op,
// and inserts built function into given module.
template <typename ClusterOrLaunchOp>

    }
  }
  OpBuilder builder(op);
  builder.setInsertionPointAfter(op);
  OperationState state(op->getLoc(), op->getName().getStringRef(),
                       op->getOperands(), result_types, op->getAttrs());
  for (int i = 0; i < op->getNumRegions(); ++i) {
    state.addRegion();
  }
  Operation* new_op = builder.create(state);

  llvm::SetVector<Value> been_defined;
  llvm::SetVector<Value> results;
  auto update_from_op = [&](Operation* op) {
    been_defined.insert(op->getResults().begin(), op->getResults().end());
    for (Value input : op->getOperands()) {
      if (been_defined.contains(input)) {
        continue;
      }
      results.insert(input);
    }
  };
  for (Operation* op : partition_ops) {
    update_from_op(op);

  // then, they are represented as mlir::TF::IfOp. We add them here, anyway, to
  // be future-proof.
  if (llvm::isa<TF::IfOp, TFL::IfOp, TFL::CallOnceOp, TFL::WhileOp>(op))
    return true;
  for (auto operand : op->getOperands()) {
    if (IsResourceTensor(operand)) return true;
  }
  for (auto result : op->getResults()) {
    if (IsResourceTensor(result)) return true;
  }
  return false;
}

  // dependency from the oprands of the op to the current cluster.
  llvm::SetVector<Operation*> ops_depend_on_cluster =
      GetAllOpsDependOnCluster(c, op_to_cluster_map);
  return llvm::none_of(to_merge->getOperands(), [&](Value value) {
    Operation* defining_op = value.getDefiningOp();
    return defining_op && ops_depend_on_cluster.contains(defining_op);
  });
}
}  // namespace

      value_select.getOperand(0).getDefiningOp());
  if (!value_gt ||
      value_gt.getComparisonDirection() != compare_direction_included ||
      value_gt.getLhs() != body.getArgument(0) ||
      value_gt.getRhs() != body.getArgument(2))
    return failure();

  mhlo::SelectOp index_select = llvm::dyn_cast_or_null<mhlo::SelectOp>(
      return_op.getOperand(1).getDefiningOp());

bool IsConvertOpLegal(ConvertOp convert_op, BFloat16TypeConverter& converter) {
  if (!converter.isLegal(convert_op.getOperand().getType())) {
    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
        convert_op.getOperand().getDefiningOp());
    return other_convert_op &&
           converter.isLegal(other_convert_op.getOperand().getType());
  } else if (!converter.isLegal(convert_op.getResult().getType())) {

bool IsHybridQuantizedOp(Operation* op) {
  if ((op->getNumOperands() != 2 && op->getNumOperands() != 3) ||
      op->getResultTypes().size() != 1) {
    return false;
  }
  Type lhs_type = op->getOperand(0).getType();
  Type rhs_type = op->getOperand(1).getType();
  Type result_type = op->getResult(0).getType();
  return !IsQuantizedTensorType(lhs_type) && IsQuantizedTensorType(rhs_type) &&
         !IsQuantizedTensorType(result_type);
}