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

      else if (auto stats_op =
                   llvm::dyn_cast_or_null<mlir::quantfork::StatisticsOp>(
                       operand.getDefiningOp()))
        operands.push_back(tensor_index_map.lookup(stats_op.getArg()));
      else
        operands.push_back(tensor_index_map.lookup(operand));
    }

    // CustomTfOp is just a wrapper around a TF op, we export the custom Op