up creating nested XLA launch ops.

    For example, the `tf.PartitionedCall` operation in the following code

    ```mlir
    func.func @convert_partitioned_call_with_resources(%arg0: tensor<!tf_type.resource>, %arg1: tensor<i32>) -> tensor<i32> {

}

void CreateSubgraphCallingFunctionWithRefVar(const Scope& s) {
  NameAttrList ref_float_function;
  ref_float_function.set_name("RefFloatFn");
  ops::PartitionedCall call(s.WithOpName("RefFloat"), {absl::Span<Input>{}},
                            {DT_FLOAT}, ref_float_function);
  Output constant =
      ops::Const(s.WithOpName("constant_ref_pco"), Input::Initializer(0.0));

// -----

// This test makes sure functions without tf._noinline=true is inlined.

module {
  func.func @partitioned_call(%arg0: tensor<1x2x2x3xf32>) -> (tensor<1x2x2x3xf32>) {
    %0 = "tf.PartitionedCall"(%arg0) <{
      config = "", config_proto = "", executor_type = "", f = @some_func
    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>

        tf_executor.graph {
          %control_0 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_1}> : () -> ()
          %control_1 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_2}> : () -> ()

  // Collect all the ops that needs to have token input names attributes. These
  // ops are communication ops and all their parent ops via nesting or function
  // calls. For example, IfRegion op and PartitionedCall op.
  std::vector<Operation*> worklist;
  absl::flat_hash_set<Operation*> ops_with_tokens;
  module.walk([&](Operation* op) {
    if (IsCommunicationOp(op)) {
      ops_with_tokens.insert(op);

  // Guarantee all functions have one use, which enables more exact shape
  // inference.
  pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  // Encapsulate PartitionedCall ops within a cluster so that the composite
  // resource ops can be decomposed.
  pm.addPass(tensorflow::tf2xla::internal::CreateXlaClusterFormationPass());

  for (int i = 0; i < num_results; ++i) {
    main_func.setResultAttr(
        i, kTfSavedModelIndexPathAttr,
        ArrayAttr::get(context, {StringAttr::get(context, output_names[i])}));
  }

  // Creates PartitionedCall ops to call exported functions.
  auto guard = OpBuilder::InsertionGuard(builder);
  int arg_idx = 0;
  int result_idx = 0;
  llvm::SmallVector<Value> call_op_returns;

}

def XlaClusterFormationPass : Pass<"tf-xla-cluster-formation", "ModuleOp"> {
  let summary = "Encapsulate partitioned calls within a Cluster op";
  let description = [{
    This pass clusters `tf.PartitionedCall` and `tf.StatefulPartitionedCall`
    with `_xla_compile_device_type` attribute into a `tf_device.cluster`.
    Notice this pass will only rewrite the outermost call if there are nested

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

      %cst = "tf.Const"() {device = "", value = dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>} : () -> tensor<2x2xf32>