%0 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsite("test@conv"("Conv2D") at "QuantizationUnit(\12\06Conv2D\1a\04conv)"))

}

// CHECK-LABEL: @float_conv_3
// CHECK:  "tf.PartitionedCall"
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_6
// CHECK:  "tf.PartitionedCall"
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_5

// CHECK-LABEL: @float_conv_1
// CHECK:  "tf.PartitionedCall"
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_2
// CHECK:  "tf.PartitionedCall"

// CHECK: %[[quantized_gather:.*]] = "tf.PartitionedCall"(
// CHECK-SAME: f = @quantized_gather_hybrid_fn_0
// CHECK: %[[float_conv:.*]] = "tf.PartitionedCall"(%[[quantized_gather]]
// CHECK-SAME: f = @composite_conv2d_fn_2
// CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%[[float_conv]]
// CHECK-SAME: f = @quantize_i8
// CHECK: %[[quantized_conv:.*]] = "tf.PartitionedCall"(%[[quantize]]

    %weight = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
    %1 = "tf.PartitionedCall"(%arg0, %weight) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>

    %1 = "tf.PartitionedCall"(%0) <{
      config = "", config_proto = "", executor_type = "", f = @some_other_func
    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    // CHECK-SKIP: tf.PartitionedCall
    // CHECK-NOSKIP: call @some_other_func
    // CHECK-NOSKIP-NOT: tf.PartitionedCall
    func.return %1: tensor<1x2x2x3xf32>
  }

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

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

// -----

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

  // 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} {