std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
    TF_CHECK_OK(s.ToGraph(g.get()));
    auto node_name_image = g->BuildNodeNameIndex();
    node_name_image["identity0"]->AddAttr("_oc", "0");
    node_name_image["identity1"]->AddAttr("_oc", "1");
    PartialTensorShape shape({2});
    node_name_image["identity1"]->AddAttr(
        kXlaInferredShapesAttrName, std::vector<PartialTensorShape>{shape});