permute¶

nvtripy.permute(input: Tensor, perm: Sequence[int]) → Tensor[source]¶

Returns a tensor with its dimensions permuted.

Parameters:

input (Tensor) – The input tensor.
perm (Sequence[int]) – The desired ordering of dimensions. It must contain all integers in \([0..N-1]\) exactly once, where \(N\) is the rank of the input tensor.

Returns:

A new tensor.

Return type:

Tensor

INPUT REQUIREMENTS:: input.dtype is one of [float32, float16, bfloat16, int4, int8, int32, int64, bool]
OUTPUT GUARANTEES:: return[0].dtype == input.dtype

Example

input = tp.reshape(tp.arange(6, dtype=tp.float32), (2, 3))
output = tp.permute(input, (1, 0))

Local Variables¶

>>> input
tensor(
    [[0, 1, 2],
     [3, 4, 5]], 
    dtype=float32, loc=gpu:0, shape=(2, 3))

>>> output
tensor(
    [[0, 3],
     [1, 4],
     [2, 5]], 
    dtype=float32, loc=gpu:0, shape=(3, 2))