Replaces operators in the graph. Args: dst: target vars representing the graph. oprmap: the map that specifies how to replace the operators. Returns: new vars that correspond to ``dst`` with all the dependencies replaced.
(dst: List[_VarNode], oprmap: Dict[_OpNode, _OpNode])
| 306 | |
| 307 | |
| 308 | def replace_oprs(dst: List[_VarNode], oprmap: Dict[_OpNode, _OpNode]) -> List[_VarNode]: |
| 309 | """Replaces operators in the graph. |
| 310 | |
| 311 | Args: |
| 312 | dst: target vars representing the graph. |
| 313 | oprmap: the map that specifies how to replace the operators. |
| 314 | |
| 315 | Returns: |
| 316 | new vars that correspond to ``dst`` with all the dependencies replaced. |
| 317 | """ |
| 318 | dst_vec = [] |
| 319 | repl_src_vec = [] |
| 320 | repl_dst_vec = [] |
| 321 | for i in dst: |
| 322 | assert isinstance(i, _VarNode) |
| 323 | dst_vec.append(i) |
| 324 | |
| 325 | for i, j in getattr(oprmap, "items", lambda: oprmap)(): |
| 326 | assert isinstance(i, _OpNode) |
| 327 | assert isinstance(j, _OpNode) |
| 328 | repl_src_vec.append(i) |
| 329 | repl_dst_vec.append(j) |
| 330 | |
| 331 | return _imperative_rt.graph._replace_oprs(repl_src_vec, repl_dst_vec, dst_vec) |
| 332 | |
| 333 | |
| 334 | def find_vars_by_name(dst: List[_VarNode], names: List[str]) -> List[_VarNode]: |