asc.experimental.asctile.transpose

asc.experimental.asctile.transpose(input: LocalTensor, *axis: int) → LocalTensor

Rearrange tensor dimensions in specific order. The supported data types are: int8, int16, int32, float16, bfloat16, float32.

Parameters:

  • input – The input tensor

  • axis – Order of input dimensions in result. Swaps two last dimensions when no axis provided

Returns:

The transposed tensor with swapped dimensions

Return type:

LocalTensor

Raises:

  • TypeError – If input is not a LocalTensor

  • RuntimeError – If the input tensor dtype is not supported or axis is incorrect

Note

If the input tensor was created by copy_in() and used only as transpose() argument, both operations will be fused into a single data copy operation during the compilation. In this case any 2D, 3D, 4D tensor is supported.

If the input is used by other operations or not created by copy_in(), a standalone transpose is used. In this case only 2D tensors in UB are supported, and input shape must be multiple of 16 (for 2 or 4 byte elements) or 32 (for 1 byte elements).

Examples

Transpose a 2D tensor:

input = asctile.copy_in(x, [0, 0, 0], [32, 16])
result = input.transpose()  # shape becomes [32, 16], same as input.transpose(1, 0)

Transpose a 3D tensor with specific order:

input = asctile.copy_in(x, [0, 0, 0], [32, 64, 16])
result = input.transpose(2, 0, 1)  # shape becomes [16, 32, 64]

Transpose as a standalone operation:

input = asctile.copy_in(x, [0, 0], [64, 64])
temp = input + 2.0  # local tensor modified after the copy_in
result = temp.transpose()

This function can also be called as a member function on LocalTensor, as input.transpose(...) instead of transpose(input, ...).