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());

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

  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.

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

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

  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;
}

  for (Operation& head_outside_compiled_op : *launch_block) {
    for (Value result : head_outside_compiled_op.getResults()) {
      bool has_external_uses = false;
      for (Operation* user : result.getUsers()) {
        if (OpInBlock(user, launch_block)) continue;
        has_external_uses = true;
        break;
      }
      if (has_external_uses) {
        launch_results.push_back(result);

        llvm::dyn_cast_or_null<quantfork::StatisticsOp>(
            op.getOperand(input_index).getDefiningOp()),
        llvm::dyn_cast_or_null<quantfork::StatisticsOp>(
            *op.getResult().getUsers().begin()),
    };

    if (!stats_ops[0] || !stats_ops[1]) {
      return failure();  // Already converted to Q-DQ pair.
    }

    llvm::SmallVector<llvm::APFloat, 4> min_max_values;

    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<quantfork::QuantizeCastOp>(connected_op);
      if (q_op && q_op.getType() == cast_type) {
        auto dq_op = llvm::cast<quantfork::DequantizeCastOp>(