%bias_scale : tensor<*xf32>, %bias_zp : tensor<*xi32>,
                           %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
        attributes {tf_quant.quantized_ops = ${quantized_ops}} {
      %0 = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
                                  %filter_scale, %filter_zp) {

                         %input : tensor<*xf32>, %weight : tensor<*xi8>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      attributes {tf_quant.quantized_ops = ${quantized_ops}} {

    %input_scale, %input_zp = "tf.PartitionedCall"(%input) {
        config = "", config_proto = "", executor_type = "", f=@internal_calculate_quant_params

  // CHECK: return

  func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<10x1x1024xf32>, %arg1: tensor<10x1024x3xf32>) -> tensor<10x1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<10x1x3xf32>

  // CHECK-LABEL: composite_conv2d_with_bias_and_relu6_fn_10
  func.func private @composite_conv2d_with_bias_and_relu6_fn_10(%arg0: tensor<1x3x2x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x2x2x2xf32> attributes {tf.tf_quant.composite_function} {
    %0 = "quantfork.stats"(%arg1) {layerStats = dense<[-3.54062747, 0.54742622]> : tensor<2xf32>} : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2xf32>

                          %bias_scale : tensor<*xf32>, %bias_zp : tensor<*xi32>,
                          %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
    attributes {tf_quant.quantized_ops = ${quantized_ops}} {
      // Given the convolution takes 2 qint8 inputs and output a qint32.
      // The accumulation scale is (input_scale * filter_scale).
      // The accumulation zero point is 0.

    return %2 : tensor<*xf32>
  }

  func.func private @composite_conv2d_with_bias_and_relu6_fn_10(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<*xf32> attributes {tf.tf_quant.composite_function} {
    %0 = "quantfork.stats"(%arg1) {layerStats = dense<[-9.54062747, 9.54742622]> : tensor<2xf32>} : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2xf32>

    %0 = "tf.PartitionedCall"(%arg0, %arg1) {config = "", config_proto = "", executor_type = "", f = @inner_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }

  func.func private @inner_fn(%arg0: tensor<1x2x2x3xf32>, %arg1: tensor<2x1024xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {

                                                          stablehlo_func_id);
  }

  // TODO - b/298966126: Currently quantizable functions are identified in TF
  // Quantizer via the tf_quant.composite_function UnitAttr attached to
  // func ops. We remove this attribute as this interferes with VHLO conversion.
  // Remove this temporary hack.
  for (auto func_op : module_op.getOps<func::FuncOp>()) {

// CHECK: }

// CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {
// CHECK-NEXT: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1)
// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true

constexpr StringRef kQuantizeFuncName = "quantize_i8";
constexpr StringRef kDequantizeFuncName = "dequantize_i8";
constexpr StringRef kAttrMapAttribute = "attr_map";
constexpr StringRef kQuantizedOpsAttribute = "tf_quant.quantized_ops";
constexpr StringRef kCompositeFuncPrefix = "composite_";
constexpr StringRef kQuantizedFuncPrefix = "quantized_";
constexpr StringRef kFloatOutputFuncSuffix = "_float_output_fn";