StrEq("composite_dot_general_fn_2"));
  EXPECT_TRUE(quantized_result.method().has_static_range_ptq());

  // Test that the non-quantized op is captured in `results`.
  const QuantizationResult& non_quantized_result = results.results(1);
  EXPECT_THAT(non_quantized_result.quantizable_unit().name(),
              StrEq("composite_dot_general_fn_1"));

  ASSERT_TRUE(module_op);

  auto composite_dot_general_fn =
      module_op->lookupSymbol<func::FuncOp>("composite_dot_general_fn_1");
  ASSERT_THAT(composite_dot_general_fn, NotNull());

  auto dot_general_op = FindOperationOfType<mlir::stablehlo::DotGeneralOp>(
      composite_dot_general_fn);
  EXPECT_TRUE(IsInLiftedFunc(dot_general_op));
}

// WholeModel-DAG: "tf.DumpTensor"(%[[matmul0_q]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "composite_dot_general_fn_1", log_dir_path = "/tmp/dumps/composite_dot_general_fn_1",...

    %4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",   device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module...

  // CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
  // CHECK: return %[[CUSTOM_AGGREGATOR_1]]
  // CHECK: }

  // CHECK: @composite_dot_general_fn_1
  // CHECK-NOT: tf_quant.composite_function

"tf.XlaCallModule"(%7, %5) <{Sout = [#tf_type.shape<2x3x1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>)...

  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
  %0 = "tf.XlaCallModule"(%arg0, %cst) {
    Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn,
    _original_entry_function = "composite_dot_general_fn",
    _quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {},
    device = "", dim_args_spec = [], disabled_checks = [],

  module {
    func.func @same_scale_after_composite() -> tensor<3x1xf32> {

      true);

  QuantizationSpec& user_provided_spec = *config.mutable_specs()->add_specs();
  user_provided_spec.mutable_matcher()->mutable_function_name()->set_regex(
      "composite_dot_general_fn_1");
  user_provided_spec.mutable_method()->mutable_no_quantization();

  // Test that the expanded `QuantizationSpec`s are populated first and then
  // user-provided specs are appended.
  //

      return %2 : tensor<?x2xf32>
    }