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.
  //

  [(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;

def LiftDotGeneral : Pat<
  (StableHLO_DotGeneralOp:$res
      $lhs, $rhs, $dot_dimension_numbers, $precision_config),
  (LiftAsTFXlaCallModule<"composite_dot_general_fn">
    (ArgumentList $lhs, $rhs),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),

      return %2 : tensor<?x2xf32>
    }

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

        disabled_checks = [], function_list = [],
        has_token_input_output = false, module = "", platforms = [],
        version = 5 : i64
    }> {
        _entry_function = @composite_dot_general_fn_1,
        _original_entry_function = "composite_dot_general_fn_1",
        _tfl_quant_trait = "fully_quantizable",
        _quantization_method = "static_range_ptq { }"
    } : (tensor<?x100352xf32>, tensor<100352x10xf32>) -> tensor<?x10xf32>

  auto target_op = FindUserOfType<T>(dynamic_broadcast_in_dim_op);
  return target_op;
}

// Gets the corresponding quantized function name from the given function name.
// Example: "composite_dot_general_fn_1" => "quantized_dot_general_fn"
std::string GetQuantizedFunctionName(const StringRef func_name) {
  return Twine(kQuantizedFuncPrefix)
      .concat(func_name.rsplit(kCompositeFuncPrefix).second)
      .str();
}

// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }

// CHECK-LABEL: private @composite_dot_general_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: return %[[DOT_GENERAL:.*]] : tensor<1x1x64xf32>
// CHECK: }

// -----

// CHECK-LABEL: @dot_general_with_bias_same_shape_fn(

    // Infer dynamic shape for XlaCallModule op.
    // Original StablehHLO function:
    //
    // func.func private @composite_dot_general_fn_1(%arg0: tensor<i32>, %arg1: tensor<?x1024xf32>, %arg2: tensor<1024x3xf32>) -> tensor<?x3xf32> attributes {_from_xla_call_module} {
    //     %0 = stablehlo.constant dense<0> : tensor<i32>