MCPcopy Index your code
hub / github.com/NVIDIA/TensorRT-LLM / DimensionExpr

Class DimensionExpr

tensorrt_llm/python_plugin.py:103–174  ·  view source on GitHub ↗

The class to wrap `trt.IDimensionExpr` to support more pythonic methods.

Source from the content-addressed store, hash-verified

101
102
103class DimensionExpr:
104 """The class to wrap `trt.IDimensionExpr` to support more pythonic methods."""
105
106 def __init__(
107 self,
108 expr: Union[trt.IDimensionExpr, int, Type[None]],
109 exprBuilder: Union[trt.IExprBuilder, Type[None]],
110 ):
111 self.exprBuilder = exprBuilder
112 self.expr = expr
113
114 @property
115 def expr(self):
116 return self._expr
117
118 @expr.setter
119 def expr(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
120 self._expr = make_expr(self.exprBuilder, expr)
121
122 def __add__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
123 expr = make_expr(self.exprBuilder, expr)
124 return expr_operation(self.expr, expr, trt.DimensionOperation.SUM, self.exprBuilder)
125
126 def __radd__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
127 return self.__add__(expr)
128
129 def __mul__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
130 expr = make_expr(self.exprBuilder, expr)
131 return expr_operation(self.expr, expr, trt.DimensionOperation.PROD, self.exprBuilder)
132
133 def __rmul__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
134 return self.__mul__(expr)
135
136 def __sub__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
137 expr = make_expr(self.exprBuilder, expr)
138 return expr_operation(self.expr, expr, trt.DimensionOperation.SUB, self.exprBuilder)
139
140 def __rsub__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
141 expr = make_expr(self.exprBuilder, expr)
142 return expr_operation(expr, self.expr, trt.DimensionOperation.SUB, self.exprBuilder)
143
144 def __eq__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
145 expr = make_expr(self.exprBuilder, expr)
146 return expr_operation(self.expr, expr, trt.DimensionOperation.EQUAL, self.exprBuilder)
147
148 def __lt__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
149 expr = make_expr(self.exprBuilder, expr)
150 return expr_operation(self.expr, expr, trt.DimensionOperation.LESS, self.exprBuilder)
151
152 def __floordiv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
153 expr = make_expr(self.exprBuilder, expr)
154 return expr_operation(self.expr, expr, trt.DimensionOperation.FLOOR_DIV, self.exprBuilder)
155
156 def __rfloordiv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):
157 expr = make_expr(self.exprBuilder, expr)
158 return expr_operation(expr, self.expr, trt.DimensionOperation.FLOOR_DIV, self.exprBuilder)
159
160 def __truediv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]):

Callers 4

expr_operationFunction · 0.85
dimsMethod · 0.85
__getitem__Method · 0.85
__setitem__Method · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected