}
      // If the result is already consumed by qint->int CastOp, manually replace
      // its use by the new UQ op. This is because such CastOp is already legal,
      // it will not go through any conversion pattern later. Without this, that
      // CastOp will still be consuming the original UQ op and cause errors.
      op->getResult(i).replaceUsesWithIf(
          new_op->getResult(i), [](OpOperand &operand) {

      }

      return true;
    }
    return false;
  }

  // Insert CastOp which is used to for converting float32 ConstantOp into
  // float16 quantization. If there is an existing CastOp connected to the
  // ConstantOp, the quantize_op will be rewired to the existing CastOp. This
  // guarantees at most one CastOp is created for float32 to float16 conversion.

      if (auto returned_op =
              returned_value.get().getDefiningOp<TFR::CastOp>()) {
        auto new_type = returned_type.clone(returned_type.getElementType()
                                                .cast<quant::QuantizedType>()
                                                .getStorageType());
        auto new_op = builder.create<TFR::CastOp>(
            returned_op->getLoc(), new_type, returned_op.getArg());

}

class RemoveRedundantCast : public OpRewritePattern<CastOp> {
  using OpRewritePattern<CastOp>::OpRewritePattern;

 public:
  LogicalResult matchAndRewrite(CastOp cast_op,
                                PatternRewriter &rewriter) const override {
    auto preceding_cast =
        llvm::dyn_cast_or_null<CastOp>(cast_op.getArg().getDefiningOp());
    if (!preceding_cast) {
      return failure();

 public:
  explicit CastBf16OpsToF32(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

 private:
  LogicalResult match(Operation* op) const override {
    if (isa<TF::CastOp, TF::ConstOp>(op) ||
        op->getName().hasTrait<OpTrait::ZeroOperands>()) {
      return failure();
    }
    for (Value input : op->getOperands()) {
      if (getElementTypeOrSelf(input).isBF16()) {

  // type.
  // TODO(fengliuai): This method is required when the operand types are not set
  // by the frontend correctly.
  Value CastToNonDerivedType(PatternRewriter& rewriter, Location loc,
                             CastOp cast_op, Type input_tfr_type) const {
    auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type);
    if (!tensor_type) return cast_op.getArg();

    auto attr_names = tensor_type.getAttrKeys();

// PadOp and CastOp.
std::optional<BlockArgumentInfo> GetInputBlockArgNum(Value input) {
  auto block_arg_num = GetBlockArgNum(input);
  if (block_arg_num.has_value()) return block_arg_num;

  Value next_input = input;
  auto pad_op = dyn_cast_or_null<TF::PadOp>(next_input.getDefiningOp());
  auto cast_op = dyn_cast_or_null<TF::CastOp>(next_input.getDefiningOp());

  auto mul_op = dyn_cast_or_null<TF::MulOp>(results[0].getDefiningOp());
  EXPECT_THAT(mul_op, NotNull());
  // Even though the preceding CastOp is foldable, it shouldn't be folded since
  // we are calling from the MulOp.
  EXPECT_TRUE(isa<TF::CastOp>(mul_op.getX().getDefiningOp()));
}

TEST_F(ConstantFoldingTest, NotFoldingArgument) {
  constexpr absl::string_view kModuleCode = R"mlir(
    module {

    for (int i = 0; i < counter_size; ++i) {
      Value word = counter[i];
      Value word_u64 = rewriter.create<CastOp>(loc, word_u64_type, word);
      Value new_word_u64 = rewriter.create<AddV2Op>(loc, word_u64, increment);
      Value new_word = rewriter.create<CastOp>(loc, word_type, new_word_u64);
      pack_args.push_back(new_word);

      Value overflow = rewriter.create<LessOp>(loc, new_word_u64, word_u64);

    auto block_shape_i64 =
        rewriter.create<CastOp>(loc, block_shape_i64_type, op.getBlockShape());

    auto paddings_i64_type = tensorflow::GetTypeFromTFTensorShape(
        paddings_type.getShape(), rewriter.getIntegerType(64));
    auto paddings_i64 =
        rewriter.create<CastOp>(loc, paddings_i64_type, op.getPaddings());

    auto pad00 = rewriter.create<ConstOp>(