How to build a tensor from one tensor of contents and another of indices?

user_13837264 0
Viewed 166

I'm trying to assemble a tensor based on the contents of two other tensors, like so:

I have a 2D tensor called A, with shape I * J, and another 2D tensor called B, with shape M * N, whose elements are indices into the 1st dimension of A. I want to obtain a 3D tensor C with shape M * N * J such that C[m,n,:] == A[B[m,n],:] for all m in [0, M) and n in [0, N).

I could do this using nested for-loops to iterate over all indices in M and N, assigning the right values to C at each one, but M and N are large so this is quite slow. I suspect there's some nicer, faster way of doing this using clever slicing or a built-in pytorch function, but I don't know what it would be. It looks a bit like somewhere one would use torch.gather(), but that requires all tensors to have the same number of dimensions. Does anyone know how this ought to be done?

EDIT: torch.index_select(input, dim, index) is almost what I want, but it won't work here because it requires that index be a 1D tensor, while my tensor of indices is 2D.

pytorch python
1 Answers

You could achieve this by flattening the first dimensions which let's you index A. A broadcast will be required to recover the final shape

>>> A[B.flatten(),:].reshape(*B.shape, A.size(-1))

Indexing with A[B.flatten(),:] is equivalent to torch.index_select(A, 0, B.flatten()).

user_63313690
Related