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

    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
        }
        float_val: 0.0
      }
    }
  }
}
node {
  name: "test_fn_call"
  op: "StatefulPartitionedCall"
  input: "const"
  attr {
    key: "Tin"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_FLOAT

  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
  }

  for (int i = 0; i < num_inputs; i++) {
    grad_outputs->push_back(grad[i]);
  }

  return scope.status();
}

REGISTER_GRADIENT_OP("PartitionedCall", PartitionedCallGrad);
REGISTER_GRADIENT_OP("StatefulPartitionedCall", PartitionedCallGrad);

}  // anonymous namespace
}  // namespace ops

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

    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
        }
        float_val: 0.0
      }
    }
  }
}
node {
  name: "test_fn_call"
  op: "StatefulPartitionedCall"
  input: "const"
  attr {
    key: "Tin"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_FLOAT

  (*f.mutable_attr())["T"].set_type(DT_FLOAT);
  auto results =
      PartitionedCall(scope_, std::initializer_list<Input>{x}, {DT_FLOAT}, f);
  RunTest(x, {}, results[0], {});

  auto stateful_results = StatefulPartitionedCall(
      scope_, std::initializer_list<Input>{x}, {DT_FLOAT}, f);
  RunTest(x, {}, stateful_results[0], {});
}

}  // namespace
}  // namespace ops

    %0 = tf_executor.graph {

  // When add_edges_to_output_of_downstream_nodes is true, the output edges of
  // the xla_launch_node's immediate downstream nodes would be attached to the
  // generated xla node. For example, if the original graph is
  // StatefulPartitionedCall{_xla_compile_id=1} -> XlaClusterOutput -> NodeA
  // The output graph of this function would look like the following when
  // add_edges_to_output_of_downstream_nodes is true:
  // XlaLaunch -> NodeA

    tf_executor.graph {
      %control = tf_executor.island {
        "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @partitioned}> : (tensor<*x!tf_type.resource>) -> ()
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }