No N-dimensional tranpose in PyTorch

Viewed 19715

PyTorch's torch.transpose function only transposes 2D inputs. Documentation is here.

On the other hand, Tensorflow's tf.transpose function allows you to transpose a tensor of N arbitrary dimensions.

Can someone please explain why PyTorch does not/cannot have N-dimension transpose functionality? Is this due to the dynamic nature of the computation graph construction in PyTorch versus Tensorflow's Define-then-Run paradigm?

2 Answers

Einops supports verbose transpositions for arbitrary number of dimensions:

from einops import rearrange
x  = torch.zeros(10, 3, 100, 100)
y  = rearrange(x, 'b c h w -> b h w c')
x2 = rearrange(y, 'b h w c -> b c h w') # inverse to the first

(and the same code works for tensorfow as well)

Related