How can I do jagged indexing for tensor?

The indexing array is jagged, how can I make it?

import torch
a = torch.randn(25, 300)
i = [[0,0],[1,1,1],[2,2,2,2]]
j = [[12,10],[13,24,26], [8,9,10,11]]
a[i,j] #error <-expected sequence of length 2 at dim 1 (got 3)

I’m not sure why you need the nested list. Aren’t the indices you really just need something like this?

i = [12,10,13,24,26,8,9,10,11]
j = [0,0,1,1,1,2,2,2,2]

If that is the case, you can just do the following:

j = torch.tensor([12,10,13,24,26,8,9,10,11]).long()
i = torch.Tensor([0,0,1,1,1,2,2,2,2]).long()
a[i,j]

And if the output size is not what you are looking for, just use

b = a[i,j].view(3,3)

Because I want to keep the jagged shape

[ [a[0,12], a[0,10]],
  [a[1,13], a[1,24], a[1,26]],
  ...
]