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/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),