}

  static bool IsPerTensor(const WeightOnlyPtq& weight_only_ptq) {
    const auto& input_quantized_types = weight_only_ptq.input_quantized_types();
    if (input_quantized_types.empty()) {
      return true;
    }
    auto weight_type = input_quantized_types.find(1);
    if (weight_type == input_quantized_types.end()) {
      return true;
    }
    return weight_type->second.has_per_tensor();
  }

    return %1 : tensor<1x3xf32>
  }

  const StaticRangePtq& srq_spec = conv_spec.method().static_range_ptq();
  ASSERT_THAT(srq_spec.input_quantized_types(), SizeIs(1));
  ASSERT_TRUE(srq_spec.input_quantized_types().contains(1));
  ASSERT_TRUE(srq_spec.input_quantized_types().at(1).has_dimension_specs());

  const QuantizedDimension& dimension_specs =
      srq_spec.input_quantized_types().at(1).dimension_specs();
  ASSERT_TRUE(dimension_specs.has_dimension());

    Sout = [#tf_type.shape<1x2x2x2>], _entry_function = @composite_conv_fn,
    _original_entry_function = "composite_conv_fn",
    _stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {per_tensor {}}}}",
    device = "", dim_args_spec = [], disabled_checks = [],
    has_token_input_output = false, module = "", platforms = [],
    version = 5 : i64

      // Represents a per-channel quantization for the operand index 1 with
      // quantization dimension of 3
      _quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
      platforms = [], version = 4 : i64
    } : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>

// -----

// Tests that basic convolution is properly quantized. It is per-channel
// quantized unless `enable-per-channel-quantized-weight=false`, according to
// `_quantization_method` with an `input_quantized_types` and explicit
// `dimension_specs`.

module attributes {tf_saved_model.semantics} {

  // 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;
}

// Represents a matching method that matches quantizable units by lifted

// weights, on the channel dimension.
//
// In textproto, the spec corresponds to:
//
// {
//   {matcher {function_name {regex: "composite_conv.*"}}}
//   {method {static_range_ptq
//     {input_quantized_types {
//       key: 1,
//       value {dimension_specs {dimension: 3}}}}
//   }}
// }
QuantizationSpec GetPtqSpecForConvolution(Method::MethodCase method_case) {
  QuantizationSpec spec{};

    // Look for quantized dimension specs for each quantized type and
    // populate `coeff_op_quant_dim`.
    for (const auto& [operand_idx, quantized_type] :
         static_range_ptq_spec.input_quantized_types()) {
      if (quantized_type.has_dimension_specs()) {
        spec.coeff_op_quant_dim[operand_idx] =
            quantized_type.dimension_specs().dimension();
      }
    }
  }
}

}  // namespace

          {
            _entry_function = @composite_conv_fn_1,
            _original_entry_function = "composite_conv_fn_1",
            _quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true},
            _tfl_quant_trait = "fully_quantizable"