Map tensor of shape [b,x,y] to [b,x+y]

1

Hi,

Is there a way in PyTorch to input a network (let’s assume a basic MLP network) a 3-dimensional tensor of shape [batch_size, x, y] and get at the output a 2-dimensional tensor of shape [batch_size, x+y]? I know I can reshape the input to [batch_size, x*y] and then the solution will be simple, but I assume x and y to be high dimensional and I want to avoid this multiplication.

Thanks!

2

Maybe it won’t fit your use case but you could use an nn.Conv1d layer accepting x input channels with a kernel size of y and returning x+y output channels. You would need to squeeze the spatial dimensions with a size of 1 but the shapes should then fit your requirements.