Method replace_vars

imperative/python/megengine/utils/network.py:378–399

r"""Replaces vars in the graph. Args: repl_dict: the map {old_var: new_var} that specifies how to replace the vars.

(self, repl_dict: Dict[VarNode, VarNode])

        assert not blacklist, "unused items in blacklist: {}".format(blacklist)

    def replace_vars(self, repl_dict: Dict[VarNode, VarNode]):
        r"""Replaces vars in the graph.

        Args:
            repl_dict: the map {old_var: new_var} that specifies how to replace the vars.
        """
        if not all([var.owner for var in repl_dict.values()]):
            self.add_dep_oprs(*list(repl_dict.values()))
        for var in self.all_vars:
            if var in repl_dict:
                repl_var = repl_dict[var]
                if repl_var is var:
                    continue
                for opnode in var.users:
                    # use method 'is' instead of 'in' to avoid
                    # compare VarNode use elemwise equal
                    assert any([var is _ for _ in opnode.inputs])
                    opnode.inputs = [repl_var if var is i else i for i in opnode.inputs]
                    if opnode not in repl_var.users:
                        repl_var.users.append(opnode)
                var.users.clear()
        self._compile()

    def replace_oprs(self, repl_dict: Dict[OpNode, OpNode]):
        r"""Replaces operators in the graph.