triton.language.trans

triton.language.trans(input: tensor, *dims, _semantic=None)

Permutes the dimensions of a tensor.

If the parameter dims is not specified, the function defaults to swapping the last two axes, thereby performing an (optionally batched) 2D transpose.

参数:

• input -- The input tensor.

• dims -- The desired ordering of dimensions. For example, (2, 1, 0) reverses the order dims in a 3D tensor.

dims can be passed as a tuple or as individual parameters:

# These are equivalent
trans(x, (2, 1, 0))
trans(x, 2, 1, 0)

permute() is equivalent to this function, except it doesn't have the special case when no permutation is specified.

This function can also be called as a member function on tensor, as x.trans(...) instead of trans(x, ...).

Example

import torch
import triton
import triton.language as tl


@triton.jit
def trans_kernel(x_ptr, out_ptr, D1: tl.constexpr, D2: tl.constexpr, D3: tl.constexpr):
    # (D1, D2, D3) -> (D3, D1, D2)
    d1 = tl.arange(0, D1)[:, None, None]
    d2 = tl.arange(0, D2)[None, :, None]
    d3 = tl.arange(0, D3)[None, None, :]
    x = tl.load(x_ptr + d1 * D2 * D3 + d2 * D3 + d3)
    y = tl.trans(x, [2, 0, 1])
    flat = tl.reshape(y, (D1 * D2 * D3, ))
    tl.store(out_ptr + tl.arange(0, D1 * D2 * D3), flat)


def test_trans():
    D1, D2, D3 = 2, 3, 4
    x = torch.zeros([D1, D2, D3], dtype=torch.float32).npu()
    out = torch.empty((D1 * D2 * D3, ), dtype=torch.float32, device="npu")
    trans_kernel[(1, )](out, x, D1, D2, D3)

    assert out.shape == (D1 * D2 * D3, ), f"Shape mismatch: expected {(D1*D2*D3,)} but got {out.shape}."
    print("Test passed.")


if __name__ == "__main__":
    test_trans()

Special Restrictions

• DataType: Ascend A2/A3 does not support fp64, fp8e4, fp8e5, uint16, uint32, uint64 (hardware limitation).

• Transpose on dimensions greater than 8 is not supported.