#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"

namespace {

using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::QuantizedType;
using ::stablehlo::quantization::WeightOnlyPtq;

// Inserts quantization parameters of weights for weight-only quantization and
// dynamic range quantization of `stablehlo.convolution` and
// `stablehlo.dot_general`.
class InsertWeightParamPass

// Checks if the `Method` attatched to the given `tf.XlaCallModule` op has
// `WeightOnlyPtq`.
bool HasWeightOnlyPtqMethod(TF::XlaCallModuleOp xla_call_module_op);

// Checks if an op is a `tf.XlaCallModule` op, contains 'conv' or 'dot_general'
// in its name and has `Method` with `WeightOnlyPtq`.
bool IsWeightOnlyQuantizableOp(const Operation& op);

// Lists the functions in a ModuleOp sorted by their names.

      signature_keys=None,
      tags=set(config.tf_saved_model.tags),
  )

  signature_def_map_serialized = _serialize_signature_def_map(signature_def_map)
  # Currently, only StaticRangePtq or WeightOnlyPtq is supported.
  # Consider merging the pipelines to address mixed algorithm models.
  if _has_quantization_method(config.specs, 'static_range_ptq'):
    pywrap_quantization.static_range_ptq(
        src_saved_model_path,

  return spec;
}

QuantizationSpec GetDefaultWeightOnlyPtqSpec() {
  QuantizationSpec spec{};
  spec.mutable_matcher()->mutable_function_name()->set_regex(
      "^.*(conv|dot_general).*");

  WeightOnlyPtq& weight_only_ptq_spec =
      *spec.mutable_method()->mutable_weight_only_ptq();
  if (auto [iter, inserted] =
          weight_only_ptq_spec.mutable_input_quantized_types()->try_emplace(1);
      inserted) {

  if (failed(pm.run(module_op))) {
    return absl::InternalError("Failed to run legalize TF to StableHLO.");
  }

  absl::StatusOr<mlir::ModuleOp> quantized_module_op;
  // Currently, only StaticRangePtq or WeightOnlyPtq is supported.
  // Consider merging the pipelines to address mixed algorithm models.
  if (HasQuantizationMethod(updated_config.specs(),
                            Method::MethodCase::kStaticRangePtq)) {

message StaticRangePtq {
  // Operand index -> QuantizedType mapping. Operands that are not specified
  // here will be quantized with best effort.
  map<int32, QuantizedType> input_quantized_types = 1;
}

message WeightOnlyPtq {
  // Operand index -> QuantizedType mapping. Operands that are not specified
  // here will be quantized with best effort.
  map<int32, QuantizedType> input_quantized_types = 1;
}

  EXPECT_THAT(spec.matcher().function_name().regex(),
              StrEq("^.*(conv|dot_general).*"));
  EXPECT_TRUE(spec.method().has_weight_only_ptq());

  const WeightOnlyPtq& weight_only_ptq_spec = spec.method().weight_only_ptq();

  EXPECT_THAT(weight_only_ptq_spec.input_quantized_types(),
              UnorderedElementsAre(Pair(
                  1, Truly([](const auto& quantized_type) {