Function make_expr

tensorrt_llm/python_plugin.py:62–86

Make a dimension expression. Parameters: exprBuilder: The trt.exprBuilder object. Using it to check whether dim has the same exprBuilder or to create trt.IDimensionExpr if necessary. dim: The input dim. Returns: A trt.IDimensionExpr object.

(
    exprBuilder: Union[trt.IExprBuilder, Type[None]],
    dim: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]],
)

def make_expr(
    exprBuilder: Union[trt.IExprBuilder, Type[None]],
    dim: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]],
) -> Union[trt.IDimensionExpr, Type[None]]:
    """Make a dimension expression.

    Parameters:
        exprBuilder: The trt.exprBuilder object. Using it to check whether dim has the same exprBuilder
            or to create trt.IDimensionExpr if necessary.
        dim: The input dim.

    Returns:
        A trt.IDimensionExpr object.
    """
    if isinstance(dim, DimensionExpr):
        assert exprBuilder == dim.exprBuilder
        return dim.expr
    elif isinstance(dim, int):
        return exprBuilder.constant(dim)
    elif dim is None:
        return None
    elif isinstance(dim, trt.IDimensionExpr):
        return dim
    else:
        raise Exception


def expr_operation(