We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← Shapes and Strides step 9 of 9
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.
Stuck?
PyTorch reference solution
Sign in to attempt this problem and reveal the reference solution.