Class HLOTensor

imperative/python/megengine/xla/rules/hlotensor.py:11–204

class HLOTensor:
    def __init__(self, tensor, shape=None, dtype=None) -> None:
        if isinstance(tensor, Sequence):
            assert len(tensor) > 0, "cannot create HLOTensor from empty sequence"
            if isinstance(tensor[0], int):
                tensor = np.array(tensor)
            else:
                assert len(tensor) == 1, f"cannot create HLOTensor from {tensor}"
                tensor = tensor[0]
        if isinstance(tensor, ir.OpResultList):
            assert len(tensor) == 1, f"cannot create HLOTensor from {tensor}"
            tensor = tensor[0]

        if isinstance(
            tensor, (int, float, np.int_, np.float16, np.float32, np.float64)
        ):
            tensor = tensor if dtype is None else np.array(tensor).astype(dtype)
            tensor = ir_utils.ir_constant(tensor)
        elif isinstance(tensor, np.ndarray):
            tensor = tensor if dtype is None else np.array(tensor).astype(dtype)
            tensor = ir_utils.ir_constant(tensor)
        else:
            assert isinstance(
                tensor, (ir.RankedTensorType, ir.BlockArgument, ir.OpResult)
            ), type(tensor)

        infered_shape = get_irnode_shape(tensor)
        infered_dtype = get_irnode_dtype(tensor)

        _check_shape(infered_shape, shape)
        _check_dtype(infered_dtype, dtype)

        self._tensor = tensor
        self._shape = infered_shape
        self._dtype = infered_dtype

    @property
    def shape(self):
        return tuple(self._shape)

    @property
    def dtype(self):
        return self._dtype

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def tensor(self):
        return self._tensor

    def __str__(self):
        return f"HLOTensor(shape={self.shape}, dtype={self.dtype})"

    def __eq__(self, rhs):
        from .elemwise import equal