Skip to content

← Shapes and Strides step 9 of 9

Easy Primitives

Move one axis, not all of them

permute takes a full ordering. Every axis, every time:

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

Which is fine for four axes and a habit that breaks in two ways.

Break one: the rank changes

Write permute(0, 2, 3, 1) for a 4-D tensor and it is simply wrong for a 3-D or 5-D one. Code that handles both ends up branching on x.dim(), or being quietly wrong for one of them.

Break two: nobody can read it

permute(0, 2, 3, 1) does not say what it is doing. You have to hold the original axis meanings in your head and apply the permutation by hand to see that it is “move channels to the end”. When the intent is a single move, spelling it as a total ordering hides it.

x.movedim(1, -1)      # move axis 1 to the end. That is the whole sentence.

movedim is exactly as free as permute: both reorder sizes and strides and return a view. This is a readability and robustness choice, not a performance one, which is why it is here rather than in the performance track.

It generalises

movedim takes lists too, and the axes not mentioned keep their relative order:

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

torch.transpose(x, a, b) swaps exactly two and leaves the rest, which is the other useful special case. Between the three, reach for the one that names the smallest thing.

Your task

def channels_last(x: torch.Tensor) -> torch.Tensor

Move axis 1 to the end, for a tensor of any rank of at least 2, without copying.

The starter hard-codes the 4-D permutation and raises on everything else.