::mlir::quant::QuantizationTestBase;

TEST_F(ConvertMlirModuleToExportedModelTest, SimpleGraphDefSet) {
  // Define a module a no-op main function.
  mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module attributes {tf_saved_model.semantics} {

      return %0 : tensor<1x3xf32>
    }
  )mlir";

  const OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kModuleWithCompositeDotGeneral);
  ASSERT_TRUE(module_op);

  // Create a pass manager with `SaveQuantizationReportInstrumentation` and
  // `QuantizeCompositeFunctionsPass`. Run the passes against `module_op`.

      return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
    }
  )mlir";

  const OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kQuantizedDotGeneral);
  ASSERT_TRUE(module_op);

  const QuantizationReport report(*module_op);
  const QuantizationResults& results = report.GetQuantizationResults();
  ASSERT_THAT(results.results(), SizeIs(1));

        tf_executor::TensorFlowExecutorDialect, quant::QuantizationDialect,
        quantfork::QuantizationForkDialect>();
  }

  // Parses `module_op_str` to create a `ModuleOp`.
  OwningOpRef<ModuleOp> ParseModuleOpString(
      const absl::string_view module_op_str) {
    return parseSourceString<ModuleOp>(module_op_str, ctx_.get());
  }

  // Convenience function that returns the first operation of type `OpT` from

TEST_F(PreCalibrationComponentTest,
       HasCustomAggregatorOpAndQuantizableFuncForSimpleDotGeneral) {
  PreCalibrationComponent component(ctx_.get());
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module attributes {} {
      func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> attributes {} {
        %0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>