Chose rows of 3d Tensor based on some repeated indices. Tricky Slicing
So this is a tricky bit of Tensor slicing I'm trying.
I have a tensor A which is 3d
&gt;&gt; A.shape
torch.Size([60, 10, 16])

So this tensor is composed of 5 different data samples, where in dim=0 we have cuts at split_ids = [10, 14, 10, 12, 14] i.e first 10 elements belong to sample1; next 14 belong to sample2 &amp; so on.. I can split the tensor in such a way:
&gt;&gt; torch.split(A, split_ids, dim=0)
(tensor([[[-0.3888, -...Backward&gt;), tensor([[[ 2.6473e-0...Backward&gt;), tensor([[[ 1.1621,  ...Backward&gt;), tensor([[[ 0.1953, -...Backward&gt;), tensor([[[ 8.1993e-0...Backward&gt;))

This comprises of tuple of 5 elements (or 5 tensors), of shapes Size(10,10,16);Size(14,10,16); and so on for the splits we had.
Now, comes the tricky part - I have another indices mapping that I have derived some previous processing for each of these individual splits. Its a list of 1d tensors like this:
&gt;&gt; reverse_map
[tensor([1, 2, 2, 1, ...='cuda:0'), tensor([ 7,  7,  9, ...='cuda:0'), tensor([7, 7, 4, 3, ...='cuda:0'), tensor([ 9,  4,  9, ...='cuda:0'), tensor([ 0,  0,  0, ...='cuda:0')]
&gt;&gt; reverse_map[0]
tensor([1, 2, 2, 1, 1, 2, 0, 1, 2, 0, 0, 1, 6, 1, 7, 1, 2, 7, 4, 5, 3, 4, 9, 7,
        7, 3, 8, 7, 7, 7], device='cuda:0')

So I need to basically use these indices and pull the count of these indices from the above split tensors i.e For tensor.Size(10,10,16) I need to pull [0, 0:3, :] which is for index 0 in dim=0, I need to pull 0:3 in dim=1 because there are 3 0's in the indexing. Then at index 1 I need to pull first 4 vectors since there are 4 1's.. and so on.
Whats the best way to do this ? Does scatter_() help here ?

I think a combination of torch.bincount and this answer can give you what you want.
For simplicity, let's focus on the first tensor split from A and the first reverse_map. You can then apply this code to the other splits and reverse_maps. Let source be the first split of shape (10, 10, 16).
Here's how it goes:
# inputs
source = torch.arange(10*10*16).view(10, 10, 16)
reverse_map = torch.tensor(tensor([1, 2, 2, 1, 1, 2, 0, 1, 2, 0, 0, 1, 6, 1, 7, 1, 2, 7, 4, 5, 3, 4, 9, 7, 7, 3, 8, 7, 7, 7])

# how many columns to pull for each row - use bincount to find out!
lengths = torch.bincount(reverse_map, minlength=source.shape[0])

# use a mask to pull the elements
mask = torch.zeros(source.shape[0], source.shape[1] + 1, dtype=source.dtype, device=source.device)
mask[(torch.arange(source.shape[0]), lengths)] = 1
mask = mask.cumsum(dim=1)[:, :-1] == 0

# expand the mask to dim=2 as well and pull the elements
out = source[mask[..., None].expand(-1, -1, source.shape[2])]

# since you pull different number of columns per row, you loose the shape of source. You need a final split to recover it
target = torch.split(out, (lengths * source.shape[2]).cpu().numpy().tolist())

target = [t_.view(-1, source.shape[2]) for t_ in target]

The output target is a list of 2-D tensors with varying number of rows (according to the counts of reverse_map and source.shape[2] columns in each.
