PyTorch: apply mapping over singleton dimension of tensor
I'm afraid the title is not very descriptive but I could not think of a better one. Essentially my problem is the following:

I have a pytorch tensor of shape (n, 1, h, w) for arbitrary integers n, h and w (in my specific case this array represents a batch of grayscale images of dimension h x w).

I also have another tensor of shape (m, 2) which maps every possible value in the first array (i.e. the first array can contain values from 0 to m - 1) to some tuple of values. I would like to "apply" this mapping to the first array so that I obtain an array of shape (n, 2, h, w).

I hope this is somewhat clear, I find this hard to express in words, here's a code example (but note that that is not super intuitive either due to the four dimensional arrays involved):

import torch

m = 18

# could also be arbitrary tensor with this shape with values between 0 and m - 1
a = torch.arange(m).reshape(2, 1, 3, 3)

# could also be arbitrary tensor with this shape
b = torch.LongTensor(
    [[11, 17, 9, 6, 5, 4, 2, 10, 3, 13, 14, 12, 7, 1, 15, 16, 8, 0],
     [11, 8, 4, 14, 13, 12, 16, 1, 5, 17, 0, 10, 7, 15, 9, 6, 2, 3]]).t()

# I probably have to do this and the permute/reshape, but how?
c = b.index_select(0, a.flatten())

# ...

# another approach that I think works (but I'm not really sure why, I found this
# more or less by trial and error). I would ideally like to find a 'nicer' way
# of doing this
c = torch.stack([
    b.index_select(0, a_.flatten()).reshape(3, 3, 2).permute(2, 0, 1)
    for a_ in a
])

# the end result should be:
#[[[[11, 17,  9],
#   [ 6,  5,  4],
#   [ 2, 10,  3]],
#
#  [[11,  8,  4],
#   [14, 13, 12],
#   [16,  1,  5]]],
#
#
# [[[13, 14, 12],
#   [ 7,  1, 15],
#   [16,  8,  0]],
#
#  [[17,  0, 10],
#   [ 7, 15,  9],
#   [ 6,  2,  3]]]]


How can I perform this transformation in an efficient manner? (Ideally not using any additional memory). In numpy this could easily be achieved with np.apply_along_axis but there seems to be no pytorch equivalent to that.

Here is one way using slicing, stacking, and view-based reshape:

In [239]: half_way = b.shape[0]//2

In [240]: upper_half = torch.stack((b[:half_way, :][:, 0], b[:half_way, :][:, 1]), dim=0).view(-1, 3, 3)
In [241]: lower_half = torch.stack((b[half_way:, :][:, 0], b[half_way:, :][:, 1]), dim=0).view(-1, 3, 3)

In [242]: torch.stack((upper_half, lower_half))
Out[242]: 
tensor([[[[11, 17,  9],
          [ 6,  5,  4],
          [ 2, 10,  3]],

         [[11,  8,  4],
          [14, 13, 12],
          [16,  1,  5]]],


        [[[13, 14, 12],
          [ 7,  1, 15],
          [16,  8,  0]],

         [[17,  0, 10],
          [ 7, 15,  9],
          [ 6,  2,  3]]]])


Some caveats are that this would work only for n=2. However, this is 1.7x faster than your loop based approach, but involves more code.



Here is a more generalized approach, which scales to any positive integer n:

In [327]: %%timeit
     ...: block_size = b.shape[0]//a.shape[0]
     ...: seq_of_tensors = [b[block_size*idx:block_size*(idx+1), :].permute(1, 0).flatten().reshape(2, 3, 3).unsqueeze(0)  for idx in range(a.shape[0])]
     ...: torch.cat(seq_of_tensors)
     ...: 
23.5 µs ± 460 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)


You can also use a view instead of reshape:

block_size = b.shape[0]//a.shape[0]
seq_of_tensors = [b[block_size*idx:block_size*(idx+1), :].permute(1, 0).flatten().view(2, 3, 3).unsqueeze(0)  for idx in range(a.shape[0])]
torch.cat(seq_of_tensors)
# outputs
tensor([[[[11, 17,  9],
          [ 6,  5,  4],
          [ 2, 10,  3]],

         [[11,  8,  4],
          [14, 13, 12],
          [16,  1,  5]]],


        [[[13, 14, 12],
          [ 7,  1, 15],
          [16,  8,  0]],

         [[17,  0, 10],
          [ 7, 15,  9],
          [ 6,  2,  3]]]])




Note: please observe that I still use a list comprehension since we've to evenly divide our tensor b to permute, flatten, reshape, unsqueeze, and then concatenate/stack along dimension 0. It's still marginally faster than my above solution.
