%0 = "tf.StatefulPartitionedCall"(%arg0) <{
      config = "", config_proto = "", executor_type = "", f = @some_func
    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    // CHECK-SKIP: tf.StatefulPartitionedCall
    // CHECK-NOSKIP: call @some_func
    // CHECK-NOSKIP-NOT: tf.StatefulPartitionedCall
    %1 = "tf.PartitionedCall"(%0) <{

    (tensor<1x10xf32> {tf_saved_model.index_path = ["output_0"]}, tensor<1x10xi64> {tf_saved_model.index_path = ["output_1"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_x:0", outputs = "StatefulPartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    %handle_0 = "tf.VarHandleOp"() {container="c", shared_name="a"} : () -> tensor<!tf_type.resource<tensor<1x10xf32>>>

func.func @nested_entry_functions() attributes {tf.entry_function = {}} {
  tf_executor.graph {
     %control = tf_executor.island wraps "tf.StatefulPartitionedCall"() {config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @func} : () -> ()
     tf_executor.fetch
  }
  func.return
}

func.func @func() attributes {tf.entry_function = {}} {

  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));

      ::llvm::cl::desc("Skip tf.ResizeBilinear and tf.ResizeNearestNeighbor")};

  Option<bool> skip_partitioned_calls_{
      *this, "skip-partitioned-calls",
      ::llvm::cl::desc(
          "Skip tf.StatefulPartitionedCall and tf.PartitionedCall")};
};

void TFToMhloPass::runOnOperation() {
  auto func = getOperation();
  MLIRContext *context = func->getContext();

  RewritePatternSet patterns(context);

//
// CHECK-LABEL:  func.func @serving_default
// CHECK-NOT:  IfrtLoadVariable
// CHECK:      "tf.VarHandleOp"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: "tf.ReadVariableOp"
// CHECK-NEXT: "tf.StatefulPartitionedCall"
// CHECK-NEXT:  return 
//
module {
  func.func @serving_default() -> tensor<*xi32> {
    %cst = "tf.Const"() <{value = dense<"some_test.txt"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>

// -----

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

bool IsNonReplicatedGraph(const Graph& graph,
                          const FunctionLibraryDefinition* function_library) {
  auto predicate = [](const Graph& graph) {
    const std::string kStatefulPartitionedCallOp = "StatefulPartitionedCall";
    for (const Node* node : graph.nodes()) {
      auto node_op = node->type_string();
      if (node_op == kStatefulPartitionedCallOp) {
        // Functions called by StatefulfulPartitionedCall ops with