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

// RUN: tf-opt %s -tf-executor-tpu-v1-island-inlining | FileCheck %s

// CHECK-NOT: tf.PartitionedCall
// CHECK-NOT: module @_tpu_v1_compat_outlined

module {
  func.func @control_input(%arg0: tensor<i1>) -> tensor<i32> {
    %0:4 = tf_executor.graph {

    %0 = "tf.PartitionedCall"(%arg0, %cst_1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> 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} {

    %4 = "quantfork.stats"(%3) {
      layerStats = dense<[-1.0, 0.8]> : tensor<2xf32>
    } : (tensor<*xf32>) -> tensor<*xf32>
    %5 = "tf.PartitionedCall"(%4, %cst_1, %cst_2) {
      _tfl_quant_trait = "fully_quantizable", config = "", config_proto = "",
      executor_type = "", f = @composite_matmul_with_bias_fn_1

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

    %0 = "quantfork.stats"(%arg0) {layerStats = dense<[1.27501142, 149.824783]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>

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

    %accum_out = "tf.PartitionedCall"(%input, %weight) {
        config = "", config_proto = "", executor_type = "", f=@${internal_func_name}
      } : (tensor<*xf32>, tensor<*xi8>) -> tensor<*xf32>

    %out = "tf.PartitionedCall"(%accum_out, %weight_scale) {
        config = "", config_proto = "", executor_type = "", f=@internal_dequantize_f32

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