Function constant

tensorrt_llm/functional.py:1184–1217

Add a constant layer. TensorRT graphs encapsulate constant values in the form of constant layers (tensorrt.IConstantLayer). That function creates such a layer from a Numpy array of values. After compilation of the network by TensorRT, those weights are stored in the serialized

(ndarray: np.ndarray,
             as_dtype: trt.DataType | None = None,
             as_shape=None)

def constant(ndarray: np.ndarray,
             as_dtype: trt.DataType | None = None,
             as_shape=None) -> Tensor:
    '''
    Add a constant layer.

    TensorRT graphs encapsulate constant values in the form of constant layers
    (tensorrt.IConstantLayer). That function creates such a layer from a Numpy
    array of values. After compilation of the network by TensorRT, those
    weights are stored in the serialized TensorRT engine.

    Parameters:
        ndarray : numpy.ndarray
            The array of values (weights) encapsulated by this constant layer.

    Returns:
        The tensor produced by the inserted layer.
    '''
    trt_dtype = np_dtype_to_trt(ndarray.dtype) if as_dtype is None else as_dtype
    trt_shape = trt.Dims(
        ndarray.shape) if as_shape is None else trt.Dims(as_shape)
    trt_count = 1
    for i in range(len(trt_shape)):
        trt_count *= trt_shape[i]
    weights = trt.Weights(trt_dtype, ndarray.ctypes.data, trt_count)
    # Prevent underlying numpy array from going out of scope
    default_net().register_ndarray(ndarray)
    layer = default_trtnet().add_constant(trt_shape, weights)
    if not default_net().strongly_typed:
        layer.set_output_type(0, trt_dtype)
    tensor = _create_tensor(layer.get_output(0), layer)
    # TODO: remove this WAR after https://nvbugs/4359151 fixed.
    set_np_weight(default_trtnet(), layer.name, ndarray)
    return tensor


# TODO: TensorRT uses sizes of the output dimensions.