// CHECK-NEXT: [[results_1:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
  %2 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NEXT: [[results_2:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor

    ```

    `_XlaMustCompile=true` in the following code

    ```mlir

```

`_XlaMustCompile=true` in the following code

```mlir

//===--------------------------------------------------------------------===//
//  tf.{|Stateful}PartitionedCall
//===--------------------------------------------------------------------===//

// Test valid tf.PartitionedCall
// CHECK-LABEL: func @testValidPartitionedCall
func.func @testValidPartitionedCall(%arg0: tensor<i32>) -> tensor<i32> {

  NameAttrList b_name_attr;
  b_name_attr.set_name("CompilableFn");
  ops::PartitionedCall b(root.WithOpName("B"), {a, a}, {DT_FLOAT}, b_name_attr);
  NameAttrList c_name_attr;
  c_name_attr.set_name("UncompilableFn");

  ops::PartitionedCall c(root.WithOpName("C"), {a}, {DT_FLOAT}, c_name_attr);
  Output d = ops::Add(root.WithOpName("D"), b.output.front(), c.output.front());

// PartitionedCall op legalization.
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @partitioned_call
func.func @partitioned_call(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK: call @pcall_func(%arg0) : (tensor<i32>) -> tensor<i32>

                                                arg_buffer_size_is_fixed(i)};
    }
  }
  UpdateFuncType(func);
}

// Holds information about a decomposed callee function for
// PartitionedCall/StatefulPartitionedCall.
struct PartitionedCallDecompositionInfo {
  bool signature_change;
  func::FuncOp decomposed_callee;
  llvm::SmallDenseMap<int64_t, int64_t> buffer_arg_to_size_arg;

func.func @serving_default(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> tensor<3xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input:0", outputs = "PartitionedCall:0"}} {
  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>

  %3 = arith.constant dense<0> : tensor<i32>

  // CHECK: PartitionedCall
  // CHECK-SAME: f = @add
  // CHECK-SAME: device = "noodle"
  %4 = "tf.Case"(%2, %arg0, %arg1) {branches = [@sub, @add], output_shapes = [#tf_type.shape<>], device = "noodle", is_stateless = false} : (tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  // CHECK: PartitionedCall
  // CHECK-SAME: f = @sub
  // CHECK-SAME: _cluster_launch = "not_ready"

  // `potential_refined_type`. Return true if the type was changed.
  bool RefineResultType(Operation* op, Value result,
                        Type potential_refined_type);

  // Infers the shape from a (Stateful)PartitionedCall operation by looking up
  // the called function and propagating the return type.
  bool InferShapeForCall(CallOpInterface call_op);

  bool InferShapeForCast(Operation* op);