// quant.qcast op in case it overrides the information from training FakeQuant
  // ops.
  func.walk([&](quantfork::QuantizeCastOp q) {
    auto input_op = q.getArg().getDefiningOp();
    if (auto stats = dyn_cast_or_null<quantfork::StatisticsOp>(input_op)) {
      q.setOperand(stats.getArg());
      if (stats.use_empty()) stats.erase();
    }
  });

  // Step 1: forward pass: propagate any value scales which are not produces

                                PatternRewriter& rewriter) const override {
    auto preceding_qcast =
        q_op.getArg().getDefiningOp<quantfork::QuantizeCastOp>();
    if (!preceding_qcast) return failure();

    auto new_qcast = rewriter.create<quantfork::QuantizeCastOp>(
        q_op.getLoc(), q_op.getType(), preceding_qcast.getArg());
    new_qcast->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
    q_op->replaceAllUsesWith(new_qcast);

        return failure();
      }
      output_types = {new_output_type};
    } else {
      output_types = {output_type.clone(elem_type.getStorageType())};
    }

    SmallVector<Value> args = {q_op.getArg(), scale, zero_point};
    FlatSymbolRefAttr func_name =
        FlatSymbolRefAttr::get(rewriter.getStringAttr(kQuantizeFuncName));

    auto quantize_call = rewriter.create<TF::PartitionedCallOp>(

              .getOutputDimensions()
              .getDefiningOp<shape::ShapeOfOp>();
      // Check if the shape come from the original conv op.
      if (!shape_of_op ||
          shape_of_op.getArg().getDefiningOp<mhlo::ConvolutionOp>() !=
              conv_op) {
        return failure();
      }
      Value new_shape_of = rewriter.create<shape::ShapeOfOp>(

void PrepareQuantizeDRQPass::removeAllStatsOp(func::FuncOp func) {
  func.walk([&](quantfork::StatisticsOp stats_op) {
    stats_op.replaceAllUsesWith(stats_op.getArg());
    stats_op.erase();
  });
}

#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize.inc"

void PrepareQuantizeDRQPass::runOnOperation() {

    auto q =
        rewriter.create<QuantizeOpT>(op.getLoc(), result_type, op.getArg());
    q->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());

    auto dq = rewriter.create<DequantizeOpT>(op.getLoc(), op.getType(), q);
    op.getResult().replaceAllUsesWith(dq);
    q.getOperation()->replaceUsesOfWith(dq, op.getArg());
    op.erase();

    return success();
  }

 private:
  int num_bits;

  let extraClassDeclaration = [{
    // Return element type of the input tensor type. Only available when the
    // input is a MLIR built-in tensor type.
    Attribute getInputElementType() {
      if (auto ty = getArg().getType().dyn_cast<TensorType>()) {
        return TypeAttr::get(ty.getElementType());
      }
      return {};
    }
  }];

  let hasCanonicalizer = 1;
}

  for (int i = 0; i < op->getNumOperands(); ++i) {
    if (auto dq_op = dyn_cast_or_null<quantfork::DequantizeCastOp>(
            op->getOperand(i).getDefiningOp())) {
      auto type =
          mlir::cast<TensorType>(dq_op.getArg().getType()).getElementType();
      if (auto per_axis_qtype =
              mlir::dyn_cast_or_null<quant::UniformQuantizedPerAxisType>(
                  QuantizedType::getQuantizedElementType(type))) {
        return true;

void PrepareDynamicRangeQuantizePass::removeAllStatsOp(func::FuncOp func) {
  func.walk([&](quantfork::StatisticsOp stats_op) {
    stats_op.replaceAllUsesWith(stats_op.getArg());
    stats_op.erase();
  });
}

void PrepareDynamicRangeQuantizePass::runOnOperation() {
  func::FuncOp func = getOperation();
  MLIRContext* ctx = func.getContext();

  for (const Value operand : same_scale_op->getOperands()) {
    auto dq_op =
        dyn_cast_or_null<quantfork::DequantizeCastOp>(operand.getDefiningOp());
    if (!dq_op) continue;

    Operation* preceding_op = dq_op.getArg().getDefiningOp();
    if (!preceding_op) continue;

    // Check whether the preceding op is a quantized composite function.
    if (isa<func::CallOp>(preceding_op)) {