EXPECT_TRUE(default_spec.method().has_static_range_ptq());

  // Test that the expansion for convolution ops is done.
  const QuantizationSpec& conv_spec = new_config.specs().specs(1);
  EXPECT_THAT(conv_spec.matcher().function_name().regex(),
              StrEq("composite_conv.*"));
  ASSERT_TRUE(conv_spec.method().has_static_range_ptq());

  const QuantizationResult& result = results.results(0);
  EXPECT_THAT(result.quantizable_unit().name(),
              StrEq("composite_dot_general_fn"));
  EXPECT_TRUE(result.method().has_static_range_ptq());
}

TEST_F(QuantizationReportTest,
       InitializeWithModuleOpWithoutQuantizationMethodAttribute) {
  // A quantized dot_general op but the `CallOp` is missing the
  // `_quantization_method` attribute.

            enable_per_channel_quantized_weight) {}

  LogicalResult match(func::FuncOp entry_func_op,
                      const Method& quantization_method) const override {
    if (!quantization_method.has_static_range_ptq()) {
      return failure();
    }
    return MatchGemmStyleOp<DotGeneralOp>(entry_func_op);
  }

  void rewrite(func::FuncOp entry_func_op, const Method& quantization_method,

void PopulateCoeffOpQuantDimIfPerChannelQuantized(
    TF::XlaCallModuleOp xla_call_module_op, OpQuantSpec& spec) {
  absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
  if (method.ok() && method->has_static_range_ptq()) {
    // TODO: b/331145946 - Use `Method` accessors.
    const StaticRangePtq& static_range_ptq_spec = method->static_range_ptq();
    // Look for quantized dimension specs for each quantized type and

  ASSERT_THAT(results.results(), SizeIs(1));
  EXPECT_THAT(results.results(0).quantizable_unit().name(),
              StrEq("composite_dot_general_fn"));
  EXPECT_TRUE(results.results(0).method().has_static_range_ptq());
}

TEST_F(SaveQuantizationReportInstrumentationTest,
       ReportNotSavedWhenNoQuantizeCompositeFunctionsPass) {
  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(

  if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
      xla_call_module_op != nullptr) {
    absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
    if (method.ok() && method->has_static_range_ptq()) return true;
  }

  TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
  return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&