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