builder_.create<arith::ConstantOp>(cst.getLoc(), cst.getValue());
          user->setOperand(operand_num, new_constant_op);
        }
      }
    }
  });
}

void QuantizationDriver::SetupAllStates() {
  for (BlockArgument arg : fn_.getArguments()) {
    args_.push_back(arg);
    Value value = arg;
    // If the argument is quantized, it should only has one user.
    if (arg.hasOneUse()) {