// CHECK-SAME: {tf_saved_model.index_path = ["end_logits"]}
// CHECK-SAME: tf.entry_function = {inputs = "serving_default_input2:0,serving_default_input1:0", outputs = "StatefulPartitionedCall:1,StatefulPartitionedCall:0"}
// CHECK-SAME: tf_saved_model.exported_names = ["serving_default"]
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 554 : i32}, tf_saved_model.semantics} {

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

    // The order of these functions is also significant.
    // CHECK: {{.*StatefulPartitionedCall.* f = @_func_non_tpu.*}}
    // CHECK-NEXT: {{.*StatefulPartitionedCall.* f = @_func_sc_forward.*}}
    // CHECK-NEXT: {{.*StatefulPartitionedCall.* f = @_func_core_tpu.*}}
    // CHECK-NEXT: {{.*StatefulPartitionedCall.* f = @_func_sc_backward.*}}
    // CHECK-NEXT: return
    // metadata ops

    // CHECK: {{.*StatefulPartitionedCall.* f = @while_cond.*}}
    // CHECK: {{.*StatefulPartitionedCall.* f = @non_tpu.*}}
    // CHECK: {{.*StatefulPartitionedCall.* f = @start_step_0.*}}
    // CHECK: {{.*StatefulPartitionedCall.* f = @while_cond.*}}
    // CHECK: {{.*StatefulPartitionedCall.* f = @non_tpu.*}}
    // CHECK: {{.*StatefulPartitionedCall.* f = @start_step_1.*}}

      // CHECK: %[[CST:.*]] = "tf.Const"
      %result = "tf.StatefulPartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @simple_callee} : () -> tensor<2xi32>

      // Multi block regions may not.
      // CHECK-NEXT: %[[CALL:.*]] = "tf.StatefulPartitionedCall"

// CHECK: tf_device.cluster
// CHECK-NEXT: tf.StatefulPartitionedCall
// CHECK-NEXT: tf_device.return
// 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) {_xla_compile_device_type = "CPU", config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @func} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

func.func @entry_func_2(%arg0: tensor<i32>) -> tensor<i32> attributes {tf_saved_model.initializer_type = ""} {

// CHECK:      node {
// CHECK-NEXT:   name: "identity"
// CHECK-NEXT:   op: "Identity"
// CHECK-NEXT:   input: "const"
// CHECK:      node {
// CHECK-NEXT:   name: "statefulpartitionedcall"
// CHECK-NEXT:   op: "StatefulPartitionedCall"
// CHECK-NEXT:   input: "const"
// CHECK-NEXT:   input: "args_1"
// CHECK:          key: "f"
// CHECK-NEXT:     value {
// CHECK-NEXT:       func {

// CHECK-LABEL: func @simple_stateful_partitioned_call
// CHECK-NOT: "tf_device.cluster"
// CHECK: "tf.StatefulPartitionedCall"
// CHECK-NEXT: "tf.Const"
func.func @simple_stateful_partitioned_call(%arg0: tensor<i32>) -> tensor<i32> {
  %0 = "tf_device.cluster"() ({
    %2 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @stateful_pcall_func} : (tensor<i32>) -> tensor<i32>

    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
        }
        float_val: 0.0
      }
    }
  }
}
node {
  name: "statefulpartitionedcall"
  op: "StatefulPartitionedCall"
  input: "const"
  input: "args_1"
  attr {
    key: "Tin"
    value {
      list {
        type: DT_FLOAT
        type: DT_RESOURCE
      }
    }
  }