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