std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});

  Node* call;
  NameAttrList f_name_attr;
  f_name_attr.set_name(fd.signature().name());
  TF_ASSERT_OK(
      NodeBuilder("B", "StatefulPartitionedCall", &root.graph()->flib_def())
          .Input(inputs)
          .Attr("Tin", {DT_RESOURCE})
          .Attr("Tout", {DT_RESOURCE})
          .Attr("f", f_name_attr)
          .Finalize(root.graph(), &call));

      NodeWith(Op("Switch"), Inputs(Out(0, xla_compile), Out(1, xla_compile)));
  auto xla_run =
      NodeWith(Op("_XlaRun"), Inputs(Out(1, predicated_compilation_key)));
  auto tf_call =
      NodeWith(Op("StatefulPartitionedCall"),
               CtrlDeps(NodeWith(Op("Identity"),
                                 Inputs(Out(0, predicated_compilation_key)))));
  auto merge = NodeWith(Op("_XlaMerge"), Inputs(Out(tf_call), Out(xla_run)));

  // CHECK-NOT: tf.VarHandleOp
  // CHECK:  "tf._TfrtGetResource"()
  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
  // CHECK: tf.StatefulPartitionedCall
  %x = "tf.StatefulPartitionedCall"(%handle) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)

// -----

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

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

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

  let summary = "Rewrites partition calls into Xla launch ops to make the attached function run on XLA.";

  let description = [{
    This pass rewrites `tf.PartitionedCall` and `tf.StatefulPartitionedCall`
    operations with `_xla_compile_device_type` attribute in a
    `tf_device.cluster` into `tf.XlaLaunch` operations. This makes the attached
    function execute with XLA. `tf.XlaLaunch` requires resource-type arguments

  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass());
  pm.addPass(mlir::createCanonicalizerPass());

  // TF -> StableHLO legalization.
  // Skip StatefulPartitionedCall to preserve aliased functions.
  mlir::odml::AddLegalizeTFToStablehloPasses(pm, /*skip_quantization_ops=*/true,
                                             /*skip_resize=*/false,

  // CHECK: return %[[SUBGRAPH_1]] : tensor<1024x3xf32>
  // CHECK: }
}

// -----

// main function contains PartitionedCall and StatefulPartitionedCall ops which
// is used to preserve aliased functions. This test make sure stablehlo ops in
// each PartitionedCall functions are lifted.

  // we don't have any evidence that choosing a stateless partitioned call helps
  // for performance.
  ops::StatefulPartitionedCall call(
      root.WithOpName("stateful_partitioned_call"), args, n->output_types(),
      func, ops::StatefulPartitionedCall::Attrs{}.ConfigProto(config_string));

  for (const Edge* e : n->in_edges()) {
    if (e->IsControlEdge()) {

//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     outputs: {
//       key  : "output_0"
//       value: {
//         name : "StatefulPartitionedCall:0"
//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     method_name: "tensorflow/serving/predict"
//   }
// }