// a sunk clone of it. This allows for reusing a sunk constant with
      // multiple uses in the region.
      llvm::DenseMap<Value, TF::ConstOp> sunk_constant;
      Region &body = cluster.getBody();
      visitUsedValuesDefinedAbove(body, [&](OpOperand *use) {
        Value constant = use->get();
        auto const_op = dyn_cast_or_null<TF::ConstOp>(constant.getDefiningOp());
        if (!const_op) return;

             llvm::SmallPtrSetImpl<Operation*>* reachable_ops,
             llvm::SmallVectorImpl<Operation*>* ops_to_visit) {
  if (auto island = llvm::dyn_cast<IslandOp>(op)) {
    mlir::visitUsedValuesDefinedAbove(
        island.getBody(), island.getBody(), [&](OpOperand* operand) {
          VisitOpOperand(graph, operand->get(), reachable_ops, ops_to_visit);
        });
  }

    Operation* op) {
  for (auto operand : op->getOperands()) {
    if (IsResource(operand)) {
      SetPotentiallyWritten(operand);
    }
  }
  visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand* operand) {
    if (IsResource(operand->get())) {
      SetPotentiallyWritten(operand->get());
    }
  });
}

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

  bool has_replicate_operands = false;
  visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand* operand) {
    if (!ancestor_of_replicate(operand->get().getParentRegion()))
      has_replicate_operands = true;
  });

  return !has_replicate_operands;
}

  // Collect all used TPUReplicatedInput ops.
  llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> replicated_input_ops;
  llvm::SmallSet<mlir::TF::TPUReplicatedInputOp, 8> seen_ops;
  mlir::visitUsedValuesDefinedAbove(
      cluster.getBody(), cluster.getBody(), [&](mlir::OpOperand* operand) {
        Operation* def = operand->get().getDefiningOp();
        if (auto ri =