%0 = "tf.StatefulPartitionedCall"(%arg0) {_XlaMustCompile = true, config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @stateful_pcall_func} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

// CHECK: _cluster_outlined_function_name = "stateful_pcall_func"
// CHECK-SAME: allow_soft_placement = true
// CHECK: tf.Const
// CHECK: tf.Add
func.func @xla_must_compile_true(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {
  %0 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", _xla_compile_device_type = "CPU", f = @stateful_pcall_func} : (tensor<i32>) -> (tensor<i32>)

    tf_device.return %2 : tensor<i32>
  }) {cluster_attr = "cluster_attr"} : () -> tensor<i32>
  %cst = "tf.Const"() {value = dense<5> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.Add"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %1 : tensor<i32>
}

func.func @stateful_pcall_func(%arg0: tensor<i32>) -> tensor<i32> {

    ```mlir
    func.func @main() -> tensor<i32> {
      %0 = "tf.StatefulPartitionedCall"() {_xla_compile_device_type = "CPU", f = @stateful_pcall_func} : () -> (tensor<i32>)
      func.return %0 : tensor<i32>
    }

    func.func @stateful_pcall_func() -> tensor<i32> {
      %0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
      func.return %0 : tensor<i32>
    }
    ```

      %0 = "tf_device.cluster"() ({
        %1 = "tf.StatefulPartitionedCall"() {_xla_compile_device_type = "CPU", f = @stateful_pcall_func} : () -> tensor<i32>
        tf_device.return %1 : tensor<i32>
      }) : () -> tensor<i32>

      func.return %0 : tensor<i32>
    }

func.func @stateful_partitioned_call(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK: call @stateful_pcall_func([[ARG]]) : (tensor<i32>) -> tensor<i32>
  %0 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @stateful_pcall_func} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

func.func @stateful_pcall_func(%arg0: tensor<i32>) -> tensor<i32> {
  func.return %arg0 : tensor<i32>
}