std::vector<QuantState>& states,
    DenseMap<Value, QuantizationDriver::QuantStateIndex>& value_to_state,
    DenseMap<QuantizationDriver::OpWithOperandIndex,
             QuantizationDriver::QuantStateIndex>& operand_states,
    DenseMap<QuantizationDriver::OpWithResultIndex,
             QuantizationDriver::QuantStateIndex>& result_states) {
  const auto [cached, inserted] = value_to_state.try_emplace(value, 0);

//
class QuantizationDriver {
 public:
  // Type alias of int used to access `states_`.
  using QuantStateIndex = int;

  // (op, operand index) pair.
  using OpWithOperandIndex = std::pair<Operation*, int>;

  // (op, result index) pair.
  using OpWithResultIndex = std::pair<Operation*, int>;

  explicit QuantizationDriver(func::FuncOp func_op, const bool is_signed,

#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/monitoring/counter.h"

//===----------------------------------------------------------------------===//
// The prepare-quantize Pass.