#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

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) {