tf_device::ClusterOp tpu_cluster = pair.second;
    if (node_to_cluster.count(node)) {
      if (node_to_cluster[node].getOperation() != tpu_cluster) {
        node->getCallableRegion()->getParentOp()->emitOpError(
            "The same function is reachable from multiple TPU Clusters.");
      }
    } else {
      node_to_cluster[node] = tpu_cluster;

  while (block_arg || cast_op_input) {
    if (block_arg) {
      // Change on device function type/shape.
      HandleFuncOp(block_arg.getOwner()->getParentOp());
      block_arg = nullptr;
      cast_op_input = nullptr;
    } else {
      auto cast_input = cast_op_input.getX();
      // Update input type.
      auto transform_result_type =

      Operation* def = input.getDefiningOp();
      DCHECK_EQ(def->getParentOp(), graph);
      try_update_current_candidate(def);
    }

    // Check island data operands.
    island.walk([&](Operation* op) {
      for (Value input : op->getOperands()) {
        Operation* def = input.getDefiningOp();
        if (!def || def->getParentOp() != graph) continue;

        try_update_current_candidate(def);
      }

          "This op has the same `shared_name` but different type with another "
          "resource op in the function");
      signalPassFailure();
      return;
    }
    op.replaceAllUsesWith(existing_resource->getParentOp()->getResults());
    ops_to_remove.push_back(&op);
  }

  // Remove op after the loop to avoid crash.
  for (Operation* op : ops_to_remove) {
    op->erase();
  }
}

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// TODO(b/158596585): Replace this with a cost model analysis.
bool IsTrivialUnaryOperation(Operation* op) {
  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}

  if (!op) return false;
  if (llvm::isa<tf_device::ReplicateOp>(op)) return false;
  if (DirectUseOfVirtualDevice(*virtual_devices, op)) return true;
  return AncestorUsesVirtualDevice(virtual_devices, op->getParentOp());
}

// Check if op or its descendant uses a key in `virtual_devices`.
bool DescendantUsesVirtualDevice(
    const std::optional<DictionaryAttr>& virtual_devices,
    Operation* operation) {

    bool has_in_island_operands = false;
    for (auto operand : sub_op.getOperands()) {
      auto defining_op = operand.getDefiningOp();
      if (!defining_op || defining_op->getParentOp() != island) continue;
      has_in_island_operands = true;

      // Remove operands from sinks.
      // We don't remove the operand if it is a stateful functional control flow

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

void TPUAnnotateDynamicShapeInputsPass::runOnOperation() {
  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {

    if (!initial_value) {
      initial_value = builder.getArrayAttr({});
    }
  }

  if (!type) return nullptr;

  auto globalOp = builder.create<ml_program::GlobalOp>(
      builder.getBlock()->getParentOp()->getLoc(), name, type, false,
      initial_value, nullptr);
  symbol_table.insert(globalOp);

  return globalOp;
}

}  // namespace

#define GEN_PASS_DEF_LOWERVARIABLEOPSTOMLPROGRAMPASS

                    llvm::SmallPtrSetImpl<Operation*>* reachable_ops,
                    llvm::SmallVectorImpl<Operation*>* ops_to_visit) {
  Operation* def = operand.getDefiningOp();
  if (def && def->getParentOp() == graph && reachable_ops->insert(def).second) {
    // Op has not been visited, add to queue to visit later.
    ops_to_visit->push_back(def);
  }
}