Class SymTensor

tensorrt_llm/python_plugin.py:231–274

The class to represent symbolic tensors. Only contains dtype and shape information for users to write their own shape/dtype inference function.

class SymTensor:
    """The class to represent symbolic tensors.

    Only contains dtype and shape information for users to write their own shape/dtype inference function.
    """

    def __init__(
        self,
        dtype: Union[torch.dtype, np.dtype, str, trt.DataType, Type[None]],
        shape: Union[ShapeExpr, Sequence[int]],
    ):
        self.dtype = dtype
        self.shape = shape

    @property
    def shape(self) -> Union[ShapeExpr, Sequence[int]]:
        return self._shape

    @shape.setter
    def shape(self, shape: Union[ShapeExpr, Sequence[int]]):
        assert isinstance(shape, (ShapeExpr, list, tuple))
        if isinstance(shape, (list, tuple)):
            for i in shape:
                assert isinstance(i, int)
        self._shape = shape

    @property
    def dtype(self) -> Union[trt.DataType, Type[None]]:
        return self._dtype

    @dtype.setter
    def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType, Type[None]]):
        if isinstance(dtype, torch.dtype):
            self._dtype = torch_dtype_to_trt(dtype)
        elif isinstance(dtype, str):
            self._dtype = str_dtype_to_trt(dtype)
        elif isinstance(dtype, np.dtype):
            self._dtype = np_dtype_to_trt(dtype)
        elif isinstance(dtype, trt.DataType):
            self._dtype = dtype
        elif dtype is None:
            self._dtype = None
        else:
            raise TypeError(f"Unsupported dtype: {dtype}")


def _convert_return_value_to_list(ret):