// expected-remark@above {{ID: 0}}
  %r1 = "tf.ReadVariableOp"(%arg0) : (tensor<!tf_type.resource>) -> tensor<1xf32>
  // expected-remark@above {{ID: 1}}
  %two = "tf.StatefulPartitionedCall"(%one, %one) {config="", config_proto="", executor_type="", f=@add} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  // expected-remark@above {{ID: 2}}

    %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @partitioned_called_func} : (tensor<20xi32>) -> tensor<*xi32>
    // CHECK: tf.StatefulPartitionedCall
    // CHECK-SAME: (tensor<20xi32>) -> tensor<20xi32>
    %1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @stateful_partitioned_call_func} : (tensor<20xi32>) -> tensor<*xi32>
    // CHECK: tf.TPUPartitionedCall

  func.return %arg0 : tensor<i32>
}

// -----

// Test valid tf.StatefulPartitionedCall
// CHECK-LABEL: func @testValidStatefulPartitionedCall
func.func @testValidStatefulPartitionedCall(%arg0: tensor<i32>) -> tensor<i32> {
  %0 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @pcall_func} : (tensor<i32>) -> (tensor<i32>)

  return false;
}

// Returns whether the current op is not supported by the TF Lite runtime.
static bool IsUnsupportedFlexOp(const std::string& op_name) {
  return op_name == "PartitionedCall" || op_name == "StatefulPartitionedCall";
}

// Create description of operation that could not be converted.
static std::string GetOpDescriptionForDebug(Operation* inst) {
  const int kLargeElementsAttr = 16;
  std::string op_str;

  %2 = "tf.If"(%0, %arg0, %arg1) {then_branch = @add, else_branch = @sub, output_shapes = [#tf_type.shape<>], device = "noodle", is_stateless = true} : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
  // CHECK: %1 = "tf.StatefulPartitionedCall"(%0, %arg1)
  // CHECK-SAME: f = @add
  // CHECK-SAME: _underscore_attr = "something"
  // CHECK-SAME: device = "noodle"