Class Tensor

tensorrt_llm/functional.py:107–602

The class to represent dense tensors. A dense tensor is named, has a shape and contains typed elements. Each dimension of a tensor can either be static or dynamic. Static dimensions are known at engine compilation by TensorRT. Dynamic dimensions can take values determined at ru

class Tensor(object):
    '''
    The class to represent dense tensors.

    A dense tensor is named, has a shape and contains typed elements. Each
    dimension of a tensor can either be static or dynamic. Static dimensions
    are known at engine compilation by TensorRT. Dynamic dimensions can take
    values determined at runtime. The tensor can be located on the host (CPU)
    or the device (GPU).
    '''

    def __init__(self,
                 name=None,
                 dtype=None,
                 shape=None,
                 dim_range=None,
                 is_network_input=True,
                 location=trt.TensorLocation.DEVICE,
                 network=None,
                 trt_tensor=None):
        '''
        Parameters:
            name : str
                The name of the tensor.

            dtype : tensorrt.DataType
                The type of the elements of the tensor. See the TensorRT
                documentation for list of supported data types.

            shape : tensorrt.Dims
                The dimensions of the tensor. In TensorRT-LLM, tensors can have
                static or dynamic dimensions (it is possible to mix static and
                dynamic dimensions).  A static dimension is known when the
                TensorRT engine is built. A dynamic dimension can be set when
                the engine is executed. Use -1 for dynamic dimensions.

            dim_range : OrderedDict
                An ordered dictionary (the positions of the elements matter)
                that associates a name and a range of values to the dimensions.
                For a static dimension, the range must be limited to a single
                value. For a dynamic dimension, the range is defined by three
                values [min, opt, max] where min and max are, respectively, the
                smallest and largest possible values of that dimension.  The
                opt value is used by TensorRT to optimize the engine for the
                most common case.

                Assume there is N optimization profiles, each item dim_range dict is ordered by:
                 (dynamic dimension name : [profile 0 (min, opt, max), profile 1 (min, opt, max), ... profile N(min, opt, max)])
                or it's following when the dimension is static (can think as min==opt==max):
                 (static dimension name : [profile 0 value, profile 1 value, ... profile N value])
                For static dimension the profile 0-N value must be same, (TODO: can it be simplified to be only 1 value?)
                And number of keys is equal to number of optimization profiles.

            is_network_input : bool
                A boolean indicating if that tensor is an input of the network.
                Inputs must be provided by the user to run the engine.

            location : tensorrt.TensorLocation