r"""Replaces vars in the graph. Args: dst: target vars representing the graph. varmap: the map that specifies how to replace the vars. Returns: new vars that correspond to ``dst`` with all the dependencies replaced.
(
dst: List[_VarNode], varmap: Dict[_VarNode, _VarNode]
)
| 278 | |
| 279 | |
| 280 | def replace_vars( |
| 281 | dst: List[_VarNode], varmap: Dict[_VarNode, _VarNode] |
| 282 | ) -> List[_VarNode]: |
| 283 | r"""Replaces vars in the graph. |
| 284 | |
| 285 | Args: |
| 286 | dst: target vars representing the graph. |
| 287 | varmap: the map that specifies how to replace the vars. |
| 288 | |
| 289 | Returns: |
| 290 | new vars that correspond to ``dst`` with all the dependencies replaced. |
| 291 | """ |
| 292 | dst_vec = [] |
| 293 | repl_src_vec = [] |
| 294 | repl_dst_vec = [] |
| 295 | for i in dst: |
| 296 | assert isinstance(i, _VarNode) |
| 297 | dst_vec.append(i) |
| 298 | |
| 299 | for i, j in getattr(varmap, "items", lambda: varmap)(): |
| 300 | assert isinstance(i, _VarNode) |
| 301 | assert isinstance(j, _VarNode) |
| 302 | repl_src_vec.append(i) |
| 303 | repl_dst_vec.append(j) |
| 304 | |
| 305 | return _imperative_rt.graph._replace_vars(repl_src_vec, repl_dst_vec, dst_vec) |
| 306 | |
| 307 | |
| 308 | def replace_oprs(dst: List[_VarNode], oprmap: Dict[_OpNode, _OpNode]) -> List[_VarNode]: |