namespace mlir {
namespace quant {

LogicalResult RemoveIdentity::matchAndRewrite(TF::IdentityOp identity,
                                              PatternRewriter &rewriter) const {
  for (Operation *user : identity->getUsers()) {
    // Replace the op with the input if output is only used by TF ops.
    // Currently this is more on the conservative side since we need to ensure
    // every consumer op to be a TF op before applying this pattern. We can

// Check if an op result value is consumed by qint -> int TF Cast OP.
bool IsQintValueQintToIntCast(Value v) {
  if (!IsIllegalType(v.getType())) {
    return true;
  }
  if (v.getUsers().empty()) {
    return false;
  }
  return llvm::all_of(v.getUsers(), [&](OpOperand operand) {
    return IsQintToIntCast(operand.getOwner());
  });
}

// Check if an op operand value is defined by int -> qint TF Cast OP.

  // a DenseHostTensor. So as long as one of the users of a captured variable is
  // a fallback op, we should keep this variable as a fallback tensor.
  for (mlir::Operation *user : arg.getUsers()) {
    if (llvm::isa<FallbackAsyncDialect>(user->getDialect())) return true;
  }
  return false;
}

void ForEachArgConsumedByFallback(
    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action) {

      return mlir::failure();
    }

    for (auto result : old_parallel_execute.getExecuteOutputs()) {
      for (mlir::Operation* user :
           llvm::make_early_inc_range(result.getUsers())) {
        if (llvm::isa<mlir::TF::TPUPartitionedOutputV2Op>(user)) {
          assert(user->use_empty());
          user->erase();
        }
      }
    }

    for (auto operand : cluster.getOperands()) {

  worklist.push(source_op);
  while (!worklist.empty()) {
    IslandOp curr_op = worklist.front();
    worklist.pop();

    if (curr_op == target_op) break;

    for (Operation* user : curr_op.getControl().getUsers()) {
      auto user_island = dyn_cast<IslandOp>(user);
      if (!user_island) continue;
      // We have labeled `user_island` before so it also must have been added to
      // `worklist` before.

    if (member.source.is<BlockArgument *>()) continue;

    Operation *op = member.source.dyn_cast<Operation *>();
    assert(op && "member operation must be not null");

    for (Operation *user : op->getUsers()) {
      // Skip users in other blocks.
      if (user->getBlock() != op->getBlock()) continue;

      // Skip users is in the `dst_root` or `src_root` clusters, if we'll merge

      return false;
    }
  } else if (!isMemoryEffectFree(candidate_op)) {
    return false;
  }

  // Don't prune if the body block argument has any other user.
  for (Operation* op : body_block_argument.getUsers()) {
    if (op != candidate_op) return false;
  }

  // Don't prune if the condition block argument has any user.
  if (!cond_block_argument.use_empty()) return false;

  // i.e. a value that escapes.
  llvm::SmallVector<Type, 4> result_types;
  for (IslandOp new_op : islands) {
    for (Value result : new_op.getOutputs()) {
      if (llvm::any_of(result.getUsers(), [&](OpOperand user) {
            return !wrapped_ops.count(user.getOwner());
          }))
        result_types.push_back(result.getType());
    }
  }

                     llvm::SmallVectorImpl<Operation*>& erase_list,
                     llvm::SmallPtrSetImpl<Operation*>& dead_ops) {
  for (mlir::Value res : op->getResults()) {
    for (Operation* user : res.getUsers()) {
      if (!dead_ops.insert(user).second) continue;
      RecursiveRemove(user, erase_list, dead_ops);
    }
  }

  erase_list.push_back(op);

  for (auto& use : op->getOpOperands()) {

      auto asset = LookupBoundInputOfType<AssetOp>(func, i, symbol_table);

      if (!asset) continue;

      auto arg = func.getArgument(i);
      bool arg_is_deletable = true;
      for (auto user : arg.getUsers()) {
        if (auto read_op =
                llvm::dyn_cast<TF::InitializeTableFromTextFileV2Op>(user)) {
          init_table_from_text_file_ops_to_erase.push_back(read_op);
        } else {