Function shape

tensorrt_llm/functional.py:2056–2096

Add an operation to create a shape tensor. The shape tensor can either be the shape of the input tensor when the parameter dim is None or a scalar (tensor of rank 0) that corresponds to the size of dim-th dimension. Parameters: input : Tensor The input tens

(input: Tensor,
          dim: Optional[int] = None,
          cast_to_dtype: Optional[Union[str, trt.DataType]] = None,
          clip_before_cast: Sequence[int] = None)

# If dim is None, return a 1-D TensorRT LLM tensor of the size
# If dim is not None, return a 0-D TensorRT LLM tensor of the dimension size
def shape(input: Tensor,
          dim: Optional[int] = None,
          cast_to_dtype: Optional[Union[str, trt.DataType]] = None,
          clip_before_cast: Sequence[int] = None) -> Tensor:
    '''
    Add an operation to create a shape tensor.

    The shape tensor can either be the shape of the input tensor when the
    parameter dim is None or a scalar (tensor of rank 0) that corresponds to
    the size of dim-th dimension.

    Parameters:
        input : Tensor
            The input tensor from which we want to extract the shape or the
            size in one dimension.

        dim : Optional[int]
            The dimension from which to extract the size. If it is None, the
            entire shape of the input tensor is returned.

    Returns:
        A tensor that contains the shape of the input tensor (if 'dim' is None)
        or the size in the dimension 'dim' of the input tensor. If 'dim' is
        'None', that tensor has the same rank as the input tensor, otherwise
        its rank is 0.
    '''
    layer = default_trtnet().add_shape(input.trt_tensor)
    res = _create_tensor(layer.get_output(0), layer)
    if cast_to_dtype is not None:
        if clip_before_cast is not None and (cast_to_dtype == 'int32'
                                             or cast_to_dtype == trt.int32):
            assert len(
                clip_before_cast
            ) == 2, f"This parameter only expects a tuple of 2 integers (lower, upper) but got {clip_before_cast}"
            res = int_clip(res, clip_before_cast[0], clip_before_cast[1])
        res = cast(res, cast_to_dtype)

    if dim is None:
        return res

    return gather(res, dim=0, indices=dim).view([])


def gather(input: Tensor, dim: int, indices: Union[Tensor, int]) -> Tensor: