const QuantState& state = quantization_driver.GetArgQuantState(arg);
    EXPECT_TRUE(isa<quant::QuantizedType>(state.params));
  }
  for (const auto& result : quantization_driver.GetResultStates()) {
    Operation* op = result.first.first;
    const int res_index = result.first.second;
    const QuantState state =
        quantization_driver.GetResultQuantState(op, res_index);

  // result.
  void Finalize();

  SmallVector<BlockArgument, 4> GetArgs() { return args_; }

  llvm::DenseMap<std::pair<mlir::Operation*, int>, int> GetResultStates() {
    return result_states_;
  }

  DenseMap<OpWithResultIndex, QuantStateIndex> result_states_;

  // Returns the state of the block argument.
  QuantState& GetArgQuantState(BlockArgument arg) {