Method _wrap

imperative/python/megengine/core/tensor/megbrain_graph.py:41–50

(self, obj)

        self._future = None

    def _wrap(self, obj):
        if type(obj) is _imperative_rt.VarNode:
            wrapper, cache = VarNode, self._var_cache
        elif type(obj) is _imperative_rt.OperatorNode:
            wrapper, cache = OpNode, self._op_cache
        else:
            raise TypeError(type(obj))
        if obj not in cache:
            cache[obj] = wrapper(obj)
        return cache[obj]

    def _set_priority_to_id(self, dest_vars):
        set_priority_to_id(_unwrap(dest_vars))