Type cast_type = quant_type.castFromExpressedType(expressed_type);

    // Insert DQ-op if it does not exist yet. Otherwise, just rewire without
    // creating a new DQ-op.
    for (auto connected_op : op->getUsers()) {
      auto q_op = llvm::dyn_cast_or_null<Q>(connected_op);
      if (q_op && q_op.getType() == cast_type) {
        auto dq_op = llvm::cast<DQ>(q_op.getResult().use_begin()->getOwner());

  // map, and (2) initialize function worklist with functions referenced from
  // device cluster ops.
  for (auto func : module.getOps<func::FuncOp>()) {
    for (auto user : symbol_map.getUsers(func)) {
      // Populate caller-callee map.
      if (func::FuncOp caller = user->getParentOfType<func::FuncOp>())
        caller_callee_map[caller].insert(func);

    for (Operation* user :
         llvm::make_early_inc_range(cluster_result.getUsers())) {
      // Check that user has no outputs that are TPUPartitionedOutputV2
      for (auto result : user->getResults()) {
        for (Operation* user : llvm::make_early_inc_range(result.getUsers())) {
          if (llvm::isa<TF::TPUPartitionedOutputV2Op>(user)) {

  for (auto& use : uses_optional.value()) {
    if (llvm::isa<TF::WhileOp>(use.getUser())) {
      if (while_op != nullptr) {
        use.getUser()->emitOpError() << "multiple users of function.";
        return LogicalResult::failure();
      } else {
        while_op = llvm::cast<TF::WhileOp>(use.getUser());
      }
    } else {
      use.getUser()->emitOpError() << "non while use of function.";
      return LogicalResult::failure();

  bool ConvertToFloat16Constant(PatternRewriter& rewriter,
                                ConstantOp op) const {
    for (Operation* connected_op : op.getResult().getUsers()) {
      ConvertOp convert_op = dyn_cast_or_null<ConvertOp>(connected_op);
      // Skip if no convert op exists.
      if (!convert_op || convert_op.getResult().use_empty()) continue;

      // Get types.

    llvm::SmallVector<Operation*, 4> candidates;
    for (Value value : values) {
      if (incoming) {
        candidates = {value.getDefiningOp()};
      } else {
        candidates.assign(value.getUsers().begin(), value.getUsers().end());
      }
      for (Operation* candidate_op : candidates) {
        if (cluster_ops.contains(candidate_op) ||
            cluster_dependent_ops.contains(candidate_op)) {

    for (Operation* user : island.getControl().getUsers()) {
      DCHECK_EQ(user->getParentOp(), graph);
      try_update_current_candidate(user);
    }

    // Check island data results.
    Block& graph_body = llvm::cast<GraphOp>(graph).GetBody();
    for (Value result : island.getOutputs()) {
      for (Operation* user : result.getUsers()) {
        Operation* def = graph_body.findAncestorOpInBlock(*user);

      // fold_all_constants_flag.
      if (!fold_all_constants_flag_) {
        if (!IsConstOrQConstInt(op)) return;
      }

      for (auto consumer : op->getResult(0).getUsers()) {
        auto consumer_call = llvm::dyn_cast_or_null<func::CallOp>(consumer);

        if (!consumer_call) continue;

        auto function_name = consumer_call.getCallee();

  LogicalResult matchAndRewrite(quantfork::DequantizeCastOp op,
                                PatternRewriter& rewriter) const override {
    SmallVector<Operation*, 4> quantizing_ops;
    auto users = op.getResult().getUsers();
    quantizing_ops.append(users.begin(), users.end());

    bool changed = false;
    // Rewrite the floating-point ops to the quantized version, by fusing
    // preceding dequantize ops and succeding quantize ops.

  if (auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(arg)) {
    if (!Conv2DInputShapeCanTransform(arg)) return std::nullopt;
    unsigned num_users =
        std::distance(block_arg.getUsers().begin(), block_arg.getUsers().end());
    BlockArgumentInfo block_arg_info = {block_arg.getArgNumber(), num_users};
    return block_arg_info;
  }
  return std::nullopt;
}