Function unbind

tensorrt_llm/functional.py:3862–3871

Removes a tensor dimension. Returns a tuple of all slices along a given dimension, already without it.

(input: Tensor, dim: int = 0)

def unbind(input: Tensor, dim: int = 0):
    '''
    Removes a tensor dimension.

    Returns a tuple of all slices along a given dimension, already without it.
    '''
    ndim = input.ndim()
    outputs = split(input, 1, dim)
    output_shape = [input.shape[i] for i in range(ndim) if i != dim]
    return [output.view(output_shape) for output in outputs]


class AllReduceStrategy(IntEnum):