Move one axis, not all of them

Suppose axis 1 represents channels and you want it last. For a 4-D tensor, either expression gives the same axis order:

x.permute(0, 2, 3, 1)
x.movedim(1, -1)

permute lists the original axes in their new order. movedim names the axis to move and its destination; the other axes keep their relative order. For this task, movedim expresses the intent directly.

Handle different ranks

A fixed four-axis permutation raises if the input has three or five axes. Moving axis 1 to the end works for any rank of at least two:

import torch

x = torch.arange(24).reshape(2, 3, 4)
y = x.movedim(1, -1)
print(list(y.shape))  # [2, 4, 3]
print(y[0, 0].tolist())  # [0, 4, 8]
print(y.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # True

For a 2-D input, axis 1 is already last, so the axis order stays the same. A dynamically constructed permute order could also support different ranks; it is simply more work to express this single move.

A view, not a data rearrangement

For strided tensors, both operations return views sharing the input's storage. They update shape and stride metadata without copying values. That does not make them cost-free, and the resulting layout can affect operations that follow.

Moving the channel axis is also different from PyTorch's torch.channels_last memory format. That format changes strides for a 4-D image tensor while retaining its logical (N, C, H, W) shape. This exercise changes the logical axis order instead.

Related operations

Use transpose(a, b) to swap two axes. Use permute when you want to specify a complete ordering. movedim can also move several axes at once: x.movedim((0, 1), (-2, -1)) moves the first two axes to the last two positions, retaining the relative order of all remaining axes.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).