3.7.3. unit_scaling.transforms.utils.replace_node_with_function

unit_scaling.transforms.utils.replace_node_with_function(graph: Graph, source: Node, target_fn: Callable[[...], Any], args: Tuple[Any, ...] | None = None, kwargs: Dict[Any, Any] | None = None, keep_type_expr: bool = True) → None[source]

Given a source node and its accompanying graph, remove the node and replace it with a new node that represents calling the target function.

Parameters:

  • graph (Graph) – the graph in which the node is present.

  • source (Node) – the node to be replaced.

  • target_fn (Callable[..., Any]) – the function to be contained in the new node.

  • args (Optional[Tuple[Any, ...]], optional) – args of the new node. Defaults to None.

  • kwargs (Optional[Dict[Any, Any]], optional) – kwargs of the new node. Defaults to None.

  • keep_type_expr (bool, optional) – retain the type expression of the removed node. Defaults to True.