tensorflow/core/tpu/graph_rewrite/combine_tpu_embedding_load_retrieve_pass.cc
// Removes edges between individual load/retrieve nodes that are added by
// tf.function.
Removes edges between individual load/retrieve nodes that are added by
tf.function.
Removes edges between individual load / retrieve nodes that are added by tf . function .
void RemoveEdgesBetweenIndividualNodes(
    Graph* graph, const std::vector<std::string>& ops,
    const absl::flat_hash_set<Node*>& nodes) {
  const absl::flat_hash_set<std::string> ops_names(ops.begin(), ops.end());
  std::vector<const Edge*> edges_to_remove;
  for (const auto* node : nodes) {
    for (const Edge* edge : node->out_edges()) {
      if (ops_names.contains(edge->dst()->op_def().name())) {
        edges_to_remove.push_back(edge);
      }
    }
  }
  for (const auto edge : edges_to_remove) {
    graph->RemoveControlEdge(edge);
  }
}
void RemoveEdgesBetweenIndividualNodes ( Graph * graph , const std :: vector < std :: string > & ops , const absl :: flat_hash_set < Node * > & nodes ) { const absl :: flat_hash_set < std :: string > ops_names ( ops . begin ( ) , ops . end ( ) ) ; std :: vector < const Edge * > edges_to_remove ; for ( const auto * node : nodes ) { for ( const Edge * edge : node -> out_edges ( ) ) { if ( ops_names . contains ( edge -> dst ( ) -> op_def ( ) . name ( ) ) ) { edges_to_remove . push_back ( edge ) ; } } } for ( const auto edge : edges_to_remove ) { graph -> RemoveControlEdge ( edge ) ; } }
Removes edges between individual load/retrieve nodes that are added by
tf.function.
Removes edges between individual load / retrieve nodes that are added by tf . function .
// We are expecting only control edges to appear in the EdgeSet // `node->out_edges()`, and thus only those are being removed. // Separate for-loop because EdgeSet does not allow mutation // while iterating.