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) {

// 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) {

LogicalResult GetControlFlowAncestors(
    Operation* op, llvm::SmallPtrSetImpl<Operation*>& control_flow_ops,
    llvm::SmallPtrSetImpl<Block*>& control_flow_blocks) {
  Block* block = op->getBlock();
  Operation* parent = block->getParentOp();
  while (block && parent && !isa<func::FuncOp>(parent)) {
    if (!IsControlFlowOp(parent))
      return op->emitOpError()
             << "expects ancestor(s) to be of ['" << IfOp::getOperationName()

    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// Checks if `op` is nested in `block`.
bool OpInBlock(Operation* op, Block* block) {
  Block* op_block = op->getBlock();
  while (op_block) {
    if (op_block == block) return true;
    if (auto* parent_op = op_block->getParentOp()) {
      op_block = parent_op->getBlock();
    } else {
      break;

  if (seen.contains(v)) {
    return failure();  // infinite loop
  }
  seen.insert(v);

  if (auto blockArg = mlir::dyn_cast<BlockArgument>(v)) {
    Operation *op = blockArg.getOwner()->getParentOp();

    // If we're in the first block, then the argument to that block is the
    // one we're looking for.
    if (auto func = cast<func::FuncOp>(op)) {
      if (blockArg.getOwner()->isEntryBlock()) {

    //
    // to
    //
    // func.func @f(...) -> ... {
    //     yield %x : foo
    // }
    // ...
    // func.call f(...) -> foo
    auto symbol_uses = f.getSymbolUses(f->getParentOp());
    if (!symbol_uses.has_value()) {
      return failure();
    }
    for (auto use : *symbol_uses) {
      Operation* caller = use.getUser();
      bool changed = false;

    Region* region = op->getParentRegion();
    regions_with_token.insert(region);

    // If the parent is not a FuncOp, then add the parent op containing a region
    // to worklist.
    Operation* parent = region->getParentOp();
    if (!isa<func::FuncOp>(parent)) {
      if (ops_with_tokens.insert(parent).second) {
        worklist.push_back(parent);
      }
      continue;
    }

  llvm::SetVector<Value> bcasts;
  cluster->walk([&](Operation* op) {
    if (op == cluster) return WalkResult::advance();
    for (auto operand : op->getOperands()) {
      Operation* scope = operand.getParentBlock()->getParentOp();
      if (scope->isProperAncestor(replicate)) {
        bcasts.insert(operand);
      }
    }
    return WalkResult::advance();
  });
  OpBuilder builder(replicate);

    // Allow inlining in regions attached to region based control flow
    // operations only if the src region is a single block region
    return isa<IfRegionOp, CaseRegionOp, WhileRegionOp>(dest->getParentOp()) &&
           llvm::hasSingleElement(*src);
  }

  // Returns true if its legal to inline a TF operation `op` into the `dest`
  // region.
  bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,

    // the expected operations with regions. Other operations can have semantics
    // that is not compatible with fusion (e.g. region compilation).
    if (!isa<func::FuncOp, IfOp, WhileOp>(contraction->getParentOp())) {
      return rewriter.notifyMatchFailure(
          contraction,
          "fused operation must be nested inside a function, If or While");
    }