I want to create a new tensor z from two tensors, say x and y with dimensions [N_samples, S, N_feats] and [N_samples, T, N_feats] respectively. The aim is to combine both tensors on the 2nd dim by mixing the elements of the 2nd dim in a specific ordering, which is stored in a variable order with dim [N_samples, U].
The ordering is different for every sample and is basically which index to extract from which tensor. It looks like this for a given sample order[0] - [x_0, x_1, y_0, x_2, y_1, ... ], where the letter indicates the tensor and the number indicates the index of the 2nd dim. So z[0] would be
z[0] = [x[0, 0, :], x[0, 1, :], y[0, 0, :], x[0, 2, :], y[0, 1, :] ... ]
How would I achieve this? I've written something that uses torch.gather that tries to do this.
x = torch.rand((2, 4, 5))
y = torch.rand((2, 3, 5))
# new ordering of second dim
# positive means take (n-1)th element from x
# negative means take (n-1)th element from y
order = [[1, 2, -1, 3, -2, 4, 3],
[1, -1, -2, 2, 3, 4, -3]]
# simple concat for gather
combined = torch.cat([x, y], dim=1)
# add a zero padding on top of combined tensor to ease gather
zero = torch.zeros_like(x)[:, 1:2]
combined = torch.cat([zero, combined], dim=1)
def _create_index_for_gather(index, offset, n_feats):
new_index = [abs(i) + offset if i < 0 else i for i in index]
# need to repeat index for each dim for torch.gather
new_index = [[x] * n_feats for x in new_index]
return new_index
_, offset, n_feats = x.shape
index_for_gather = [_create_index_for_gather(i, offset, n_feats) for i in order]
z = combined.gather(dim=1, index=torch.tensor(index_for_gather))
Is there a more efficient way of doing this?