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"));

        # function name matches "composite_dot_general_fn_1".
        # "composite_dot_general_fn_2" will be quantized.
        specs=qc.QuantizationSpecs(
            specs=[
                qc.QuantizationSpec(
                    matcher=qc.MatcherSpec(
                        function_name=qc.FunctionNameMatcherSpec(
                            regex='composite_dot_general_fn_1'
                        )

func.func private @quantize_simple_xla_call_module_no_operand() -> tensor<1x3xf32> {

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

      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _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...

    %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>)...

  return %2 : tensor<1x3xf32>
}

// CHECK: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[CST]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1"