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,

  for (const auto& arg : quantization_driver.GetArgs()) {
    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 =

    ],
    licenses = ["notice"],
)

cc_library(
    name = "quantization_lib",
    srcs = [
        "quantization_driver.cc",
        "quantization_interface.cc.inc",
        "quantization_utils.cc",
    ],
    hdrs = [
        "quantization_driver.h",
        "quantization_interface.h.inc",
        "quantization_traits.h",
        "quantization_utils.h",
    ],
    deps = [

#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.h"

#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.

include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// Converts tf.Const to arith.constant for statically shaped, non-opaque constants.
// Needed for QuantizationDriver to recognize constants.
def ConvertTfConstToArithConst : Pat<
  (TF_ConstOp:$res DenseElementsAttr:$value),
  (Arith_ConstantOp $value),

include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// Converts tf.Const to arith.constant for statically shaped, non-opaque constants.
// Needed for QuantizationDriver to recognize constants.
def ConvertTfConstToArithConst : Pat<
  (TF_ConstOp:$res DenseElementsAttr:$value),
  (Arith_ConstantOp $value),