Function cast

tensorrt_llm/functional.py:876–919

Add a cast operation. For an input tensor of type INT8, this function sets the dynamic range of the input to [-127, 127] for automatic dequantization. For a cast into INT8, that function sets the dynamic range of the output to [-127, 127] for automatic quantization. Parame

(input: Tensor, dtype: Union[str, trt.DataType])

def cast(input: Tensor, dtype: Union[str, trt.DataType]):
    '''
    Add a cast operation.

    For an input tensor of type INT8, this function sets the dynamic range of
    the input to [-127, 127] for automatic dequantization. For a cast into
    INT8, that function sets the dynamic range of the output to [-127, 127] for
    automatic quantization.

    Parameters:
        input : Tensor
            The input tensor on which the cast is applied.

        dtype : str or trt.DataType
            The data type of the output tensor after the cast. When 'dtype' is
            provided as a string, it must be a name amongst the valid names.
            See _str_to_trt_dtype_dict in _utils.py for a list of supported
            types and type names.

    Returns:
        The tensor produced by the inserted layer.
    '''
    if isinstance(dtype, str):
        cvt_dtype = str_dtype_to_trt(dtype)
    elif isinstance(dtype, trt.DataType):
        cvt_dtype = dtype
    else:
        raise TypeError("%s is not supported" % type(dtype))

    if input.dtype == cvt_dtype:
        # If input type and cast dtype are the same, do nothing
        return input

    layer = default_trtnet().add_cast(input.trt_tensor, cvt_dtype)
    if not default_net().strongly_typed:
        layer.set_output_type(0, cvt_dtype)
    output = _create_tensor(layer.get_output(0), layer)
    if not default_net().strongly_typed:
        if input.dtype == str_dtype_to_trt('int8'):
            layer.get_input(0).set_dynamic_range(-127, 127)
        if cvt_dtype == str_dtype_to_trt('int8'):
            layer.get_output(0).set_dynamic_range(-127, 127)

    return output


def flip(input: Tensor, dims: Sequence[int]) -> Tensor: