Class ShapeExpr

tensorrt_llm/python_plugin.py:177–228

The class to Wrap `trt.DimsExprs` to support more pythonic methods.

class ShapeExpr:
    """The class to Wrap `trt.DimsExprs` to support more pythonic methods."""

    def __init__(
        self,
        dims: Union[Sequence[trt.IDimensionExpr], Sequence[int], Sequence[type[None]]],
        exprBuilder: Union[trt.IExprBuilder, type[None]],
    ):
        self.exprBuilder = exprBuilder
        self.dims = dims

    @property
    def dims(self):
        return self._dims

    @dims.setter
    def dims(
        self,
        dims: Sequence[Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]],
    ):
        if dims is not None:
            self._dims = [
                DimensionExpr(make_expr(self.exprBuilder, i), self.exprBuilder) for i in dims
            ]
        else:
            self._dims = None

    def __getitem__(self, index: int):
        if self._dims is not None:
            return self._dims[index]
        else:
            return DimensionExpr(None, self.exprBuilder)

    def __setitem__(
        self,
        index: int,
        value: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]],
    ):
        if self._dims is None:
            return
        assert index < len(self._dims)
        value = DimensionExpr(make_expr(self.exprBuilder, value), self.exprBuilder)
        self._dims[index] = value

    def __len__(self):
        if self._dims is None:
            return 0
        else:
            return len(self._dims)

    def to_trt(self) -> trt.DimsExprs:
        return trt.DimsExprs([i.expr for i in self.dims])


class SymTensor: