%i8_max = "tf.Const"() {value = dense<127> : tensor<i8>} : () -> tensor<i8>

    %clip_min = "tf.PartitionedCall"(%i8_min, %out_scale, %out_zp) {
        config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
      } : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
    %clip_max = "tf.PartitionedCall"(%i8_max, %out_scale, %out_zp) {

    %val = "tf.PartitionedCall"(%handle) {config = "", config_proto = "", executor_type = "", f = @f_callee} : (tensor<!tf_type.resource<tensor<0xf32>>>) -> (tensor<0xf32>)
    func.return %val : tensor<0xf32>
  }

  // CHECK: func private @f_callee() -> tensor<0xf32>
  func.func private @f_callee(%arg0: tensor<*x!tf_type.resource>) -> tensor<0xf32> {

    %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @serving_default} : (tensor<?x!tf_type.string>) -> (tensor<*xi64>)
    %1 = "tf.Identity"(%0) : (tensor<*xi64>) -> tensor<*xi64>
    return %1 : tensor<*xi64>
  }

module {
  func.func @wrap_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
    %0 = "tf.PartitionedCall"(%arg0, %arg1) <{f = @composite_conv2d_with_relu6_fn}> {_tfl_quant_trait = "fully_quantizable"}
          : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

      %accum_zp_f32 = "tf.ZerosLike"(%accum_scale) : (tensor<*xf32>) -> tensor<*xf32>
      %accum_zp = "tf.Cast"(%accum_zp_f32) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi32>
      %main_out = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
                                  %filter_scale, %filter_zp, %accum_scale, %accum_zp) {
          config = "", config_proto = "", executor_type = "", f=@GenerateImplFunctionName(${main_op})

  // CHECK: return %[[SUBGRAPH_1]] : tensor<1024x3xf32>
  // CHECK: }
}

// -----

// main function contains PartitionedCall and StatefulPartitionedCall ops which
// is used to preserve aliased functions. This test make sure stablehlo ops in
// each PartitionedCall functions are lifted.

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {

func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {

      attributes {tf_quant.quantized_ops = ${quantized_ops}} {

    %input_scale, %input_zp = "tf.PartitionedCall"(%input) {
        config = "", config_proto = "", executor_type = "", f=@internal_calculate_quant_params
      } : (tensor<*xf32>) -> (tensor<*xf32>, tensor<*xi32>)

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

// CHECK-NEXT:    %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK-NEXT:    %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>

    : (tensor<!tf_type.resource>, tensor<i1>) -> tensor<!tf_type.resource>
  // CHECK: "tf.PartitionedCall"
  // CHECK-SAME: f = @callee_stack_decomposed
  %call2 = "tf.PartitionedCall"(%stack, %arg0) {f = @callee, config = "", config_proto = "", executor_type = ""}
    : (tensor<!tf_type.resource>, tensor<i1>) -> tensor<!tf_type.resource>
  // CHECK: "tf.Slice"