| 1552 | } |
| 1553 | |
| 1554 | Status MutableGraphView::DeleteNodes( |
| 1555 | const absl::flat_hash_set<string>& nodes_to_delete) { |
| 1556 | TF_RETURN_IF_ERROR(CheckNodesCanBeDeleted(nodes_to_delete)); |
| 1557 | |
| 1558 | // Find nodes in internal state and delete. |
| 1559 | for (const string& node_name_to_delete : nodes_to_delete) { |
| 1560 | NodeDef* node = GetNode(node_name_to_delete); |
| 1561 | if (node != nullptr) { |
| 1562 | RemoveFaninsInternal(node, /*keep_controlling_fanins=*/false); |
| 1563 | RemoveFanoutsInternal(node); |
| 1564 | } |
| 1565 | } |
| 1566 | for (const string& node_name_to_delete : nodes_to_delete) { |
| 1567 | nodes().erase(node_name_to_delete); |
| 1568 | } |
| 1569 | |
| 1570 | // Find nodes in graph and delete by partitioning into nodes to retain and |
| 1571 | // nodes to delete based on input set of nodes to delete by name. |
| 1572 | // TODO(lyandy): Use a node name->idx hashmap if this is a performance |
| 1573 | // bottleneck. |
| 1574 | int pos = 0; |
| 1575 | const int last_idx = graph()->node_size() - 1; |
| 1576 | int last_pos = last_idx; |
| 1577 | while (pos <= last_pos) { |
| 1578 | if (nodes_to_delete.contains(graph()->node(pos).name())) { |
| 1579 | graph()->mutable_node()->SwapElements(pos, last_pos); |
| 1580 | --last_pos; |
| 1581 | } else { |
| 1582 | ++pos; |
| 1583 | } |
| 1584 | } |
| 1585 | if (last_pos < last_idx) { |
| 1586 | graph()->mutable_node()->DeleteSubrange(last_pos + 1, last_idx - last_pos); |
| 1587 | } |
| 1588 | |
| 1589 | return Status::OK(); |
| 1590 | } |
| 1591 | |
| 1592 | void MutableGraphView::RemoveFaninsInternal(NodeDef* deleted_node, |
| 1593 | bool keep_controlling_fanins) { |