Function flip

tensorrt_llm/functional.py:922–963

Reverses the order of an n-D tensor along given axis in dims. That flip operation maps to a TensorRT ISliceLayer. For the dimensions listed in dims it copies the elements from the last one to the first one (from (N-1) down to 0 with a step of -1). For the dimensions not in 'dims',

(input: Tensor, dims: Sequence[int])

def flip(input: Tensor, dims: Sequence[int]) -> Tensor:
    '''
    Reverses the order of an n-D tensor along given axis in dims.

    That flip operation maps to a TensorRT ISliceLayer. For the dimensions
    listed in dims it copies the elements from the last one to the first one
    (from (N-1) down to 0 with a step of -1). For the dimensions not in 'dims',
    it copies the elements from the first one to the last one (from 0 to N-1
    with a step of 1).

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

        dims : list or tuple
            The axes to flip. Negative indices are supported.

    Returns:
        The tensor produced by the inserted layer.
    '''
    assert not input.is_dynamic()

    ndim = input.ndim()

    for index, value in enumerate(dims):
        assert -ndim <= value < ndim
        if -ndim <= value < 0:
            dims[index] += ndim

    assert len(dims) == len(set(dims))

    start_values = [
        input.size()[i] - 1 if i in dims else 0 for i in range(ndim)
    ]
    stride_values = [-1 if i in dims else 1 for i in range(ndim)]

    layer = default_trtnet().add_slice(input.trt_tensor,
                                       start=start_values,
                                       shape=input.size(),
                                       stride=stride_values)

    return _create_tensor(layer.get_output(0), layer)


def interpolate(input: Tensor,