Class TensorNode

imperative/python/megengine/traced_module/node.py:209–305

r"""``TensorNode`` represents the Tensor objects.

class TensorNode(Node):
    r"""``TensorNode`` represents the Tensor objects."""

    _shape = None  # type: Tuple[int]
    _dtype = None  # type: numpy.dtype
    _qparams = None  # type: QParams
    _device = None
    _value = None  # type: Tensor

    def __init__(
        self,
        expr,
        name: str = None,
        qualname: str = None,
        shape: Tuple[int] = None,
        dtype: numpy.dtype = None,
        qparams: QParams = None,
    ):
        super().__init__(expr, name, qualname)
        self._shape = shape
        self._dtype = dtype
        self._qparams = qparams

    def __getstate__(self):
        state = {
            "expr": self.expr,
            "users": self.users,
            "_id": self._id,
            "_qparams": self._qparams,
            "_shape": self._shape,
            "_dtype": self._dtype,
            "_device": self._device,
            "_name": self._name,
            "_qualname": self._qualname,
        }
        _check_obj_attr(state)
        return state

    def __setstate__(self, state):
        if "_orig_name" in state:
            qualname = state.pop("_orig_name")
            modulepath, comma, qualname = qualname.rpartition(".")
            expr_name = state["expr"].__class__.__name__
            if expr_name not in ["GetAttr"]:
                qualname = "[{}]".format(qualname)
            if comma:
                qualname = "{}.{}".format(modulepath, qualname)
            state["_qualname"] = qualname
        self.__dict__.update(state)

    @property
    def shape(self):
        r"""Get the shape of this Node."""
        return self._shape

    @shape.setter
    def shape(self, shape):
        self._shape = shape