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 = {}} {

  func.func @main(%arg0: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["serving_default_input_vocabs:0"]}) -> (tensor<*xi64> {tf_saved_model.index_path = ["StatefulPartitionedCall:0"]})
      attributes {tf.entry_function = {inputs = "serving_default_input_vocabs:0", outputs = "StatefulPartitionedCall:0"}, tf_saved_model.exported_names = ["main"]} {
    %0 = tf_executor.graph {

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

  func.func @f(%arg0: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @v})
  attributes {tf_saved_model.exported_names = ["f"]} {
    "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f_callee} : (tensor<!tf_type.resource<tensor<f32>>>) -> ()
    func.return
  }

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

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

// -----

// Tests PartitionedCall/StatefulPartitionedCall.

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i1>) -> () {
  %max_size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK-NOT: tf.Stack
  %stack = "tf.StackV2"(%max_size) {elem_type = f32, stack_name = "s"} : (tensor<i32>) -> tensor<!tf_type.resource>
  // CHECK: "tf.StatefulPartitionedCall"
  // CHECK-SAME: f = @callee_stack_decomposed

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

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