using ::stablehlo::quantization::PresetQuantizationMethod;
using ::stablehlo::quantization::QuantizationComponentSpec;
using ::stablehlo::quantization::QuantizationOptions;
using QuantizationComponent =
    ::stablehlo::quantization::QuantizationComponentSpec_QuantizationComponent;
using BitType = ::stablehlo::quantization::QuantizationComponentSpec_BitType;

// Components whose QuantizationComponentSpec is not specified will not be
// quantized, and remain f32.
// NEXT ID: 7
message QuantizationComponentSpec {
  // NEXT ID: 4
  enum QuantizationComponent {
    COMPONENT_UNSPECIFIED = 0;
    COMPONENT_ACTIVATION = 1;
    COMPONENT_WEIGHT = 2;
    COMPONENT_BIAS = 3;
  }

  // NEXT ID: 5
  enum BitWidth {
    BIT_WIDTH_UNSPECIFIED = 0;

// Defines tensor type of the component. If the combination is not supported,
// an error will be raised.
// NEXT ID: 3
message QuantizationComponentSpec {
  // NEXT ID: 4
  enum QuantizationComponent {
    COMPONENT_UNSPECIFIED = 0;
    COMPONENT_ACTIVATION = 1;
    COMPONENT_WEIGHT = 2;
    COMPONENT_BIAS = 3;
  }

  // NEXT ID: 4
  enum TensorType {
    TENSORTYPE_UNSPECIFIED = 0;

    }

    if (quant_options_.quantization_method()
            .quantization_component_specs()
            .empty()) {
      auto add_new_spec =
          [this](QuantizationComponentSpec::QuantizationComponent component,
                 QuantizationComponentSpec::TensorType type) {
            QuantizationComponentSpec* new_spec =
                quant_options_.mutable_quantization_method()

_PresetMethod = _QuantizationMethod.PresetMethod
_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)

_QuantizationComponent = _QuantizationComponentSpec.QuantizationComponent
_TensorType = _QuantizationComponentSpec.TensorType

_RepresentativeDatasetFile = quant_opts_pb2.RepresentativeDatasetFile

# Mapping of signature def key -> SignatureDef.

_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)

_QuantizationComponent = (
    quant_opts_pb2.QuantizationComponentSpec.QuantizationComponent
)
_TensorType = quant_opts_pb2.QuantizationComponentSpec.TensorType

_TensorShape = Sequence[Union[int, None]]

_PER_CHANNEL_QUANTIZED_OPS = (
    'UniformQuantizedConvolution',