3.6.4. unit_scaling.transforms.prune_selected_nodes

unit_scaling.transforms.prune_selected_nodes(graph: Graph, targets: Iterable[Callable[[...], Any] | str]) → Graph

Given an FX Graph, prunes all nodes with functions in the set of target nodes.

Parameters:

• graph (Graph) – the FX graph to prune.

• targets (Iterable[Target]) – the functions which will not be represented by nodes once pruning is complete.

Returns:

the pruned graph.

Return type:

Graph