ksize = [1, 2, 2, 1], padding = "VALID", strides = [1, 2, 2, 1]
    } : (tensor<*xf32>) -> tensor<*xf32>
    %3 = "tf.Reshape"(%2, %cst) {device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
    %4 = "tf.PartitionedCall"(%3, %cst_1, %cst_2) {
      _tfl_quant_trait = "fully_quantizable", config = "", config_proto = "",
      executor_type = "", f = @composite_matmul_with_bias_fn_1

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

    %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

    %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

    %cst = "tf.Const"() {device = "", value = dense<[[[[-0.315365672, 0.27481091], [0.0901821703, -0.382271349], [-0.105572946, -0.354302853]], [[-0.47703138, -0.307006568],...

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 29 : i32}} {
  func.func @main() {
    tf_executor.graph {
      %0 = tf_executor.island wraps "tf.PartitionedCall"() {Tin = [], Tout = [], config = "", config_proto = "", device = "", executor_type = "", f = @foo, name = "Call_foo"} : () -> ()
      tf_executor.fetch
    }
    func.return
  }
  func.func @foo() {

    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<?x?xf32>

    // CHECK: %[[RESULT:.*]] = mhlo.constant dense<[10, 19]>
    %2 = "tf.PartitionedCall"(%1) {config = "", config_proto = "", executor_type = "", f = @get_shape} : (tensor<?x?xf32>) -> (tensor<?xi64>)

    // CHECK: return %[[RESULT]]
    func.return %2 : tensor<?xi64>
  }

  // CHECK: _xla_original_oc_node_name = [[NODE_NAME1:.*]], _xla_token_input_nodes = ["_xla_token_arg_node"]
  %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @PartitionedCall3} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NOT: _xla_token_input_nodes
  %1 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @IdentityFunc} : (tensor<i32>) -> (tensor<i32>)