// 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

                                                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;