}
      // 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.

}

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 {

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

    builder.setInsertionPoint(op);
    for (auto& op_operand : op->getOpOperands()) {
      if (IsLargeFloatType(op_operand.get().getType())) {
        op_operand.set(builder.create<TF::CastOp>(
            op->getLoc(), ToBfloat16Type(op_operand.get().getType()),
            op_operand.get()));
      }
    }
    builder.setInsertionPointAfter(op);
    for (auto op_result : op->getOpResults()) {

          op.getLoc(), globalOp.getType(),
          SymbolRefAttr::get(op->getContext(), globalOp.getSymName()));
      if (globalOp.getType() != op.getValue().getType()) {
        load = builder.create<TF::CastOp>(op.getLoc(), op.getValue().getType(),
                                          load->getResult(0));
      }
      op.getResult().replaceAllUsesWith(load->getResult(0));
      op.erase();
    });