Function expr_operation

tensorrt_llm/python_plugin.py:89–100

The function to do expr operation with None support.

(
    a: Union[trt.IDimensionExpr, Type[None]],
    b: Union[trt.IDimensionExpr, Type[None]],
    operation: trt.DimensionOperation,
    exprBuilder: trt.IExprBuilder,
)

def expr_operation(
    a: Union[trt.IDimensionExpr, Type[None]],
    b: Union[trt.IDimensionExpr, Type[None]],
    operation: trt.DimensionOperation,
    exprBuilder: trt.IExprBuilder,
):
    """The function to do expr operation with None support."""
    if exprBuilder is None or a is None or b is None:
        expr = None
    else:
        expr = exprBuilder.operation(operation, a, b)
    return DimensionExpr(expr, exprBuilder)


class DimensionExpr: