Output identity = ops::Identity(s.WithOpName("identity"), add);
  Graph g(OpRegistry::Global());
  TF_CHECK_OK(s.ToGraph(&g));

  TF_CHECK_OK(PerformStaticShapeInferenceBeforeEncapsulation(&g));

  // Check that "add" node now has _xla_inferred_shapes attr.
  auto node_index = g.BuildNodeNameIndex();
  Node *add_node = node_index["add"];
  std::vector<PartialTensorShape> output_shapes;

// encapsulation, some nodes will be encapsulated into function call, and shape
// inference does not handle function call at the moment.
Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g);

// Attribute indicating that some ops in this node's XLA computation has control
// dependency on this node. Attribute value will always be "true".

const char kXlaOutsideCompilationInputsAttrName[] = "_xla_oc_inputs";
const char kXlaIsPlaceholderForArg[] = "_xla_is_placeholder_for_arg";

Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g) {
  // Perform shape inference.
  std::map<int, InferredShape> arg_shapes;
  GraphShapeInfo shape_info;
  TF_RETURN_IF_ERROR(

  options.allow_internal_ops = true;
  std::unique_ptr<Graph> graph(new Graph(lib_def.get()));
  s = ConvertGraphDefToGraph(options, *graphdef, graph.get());
  if (!s.ok()) return s;

  s = PerformStaticShapeInferenceBeforeEncapsulation(graph.get());
  if (!s.ok()) return s;

  // Create FunctionLibraryRuntime.
  SessionOptions session_options;
  std::vector<std::unique_ptr<Device>> devices;