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".