Selecting from tensor

1

Hi,
I have a 3d tensor sized (batch_size, S, 5).
I also have a list containing a list for each batch, that contains indices.
For example, if the batch size is 2, the list looks like this:
[[idx0_0, idx0_1, …], [idx1_0, idx1_1, …]] (where each idx is between 0 to S).

I want to get a tensor that contains, for each batch, the 3rd dimension’s stuff from the original tensor, only from the indices that appear in the list of the same batch, while keeping the gradient of the original tensor.

For example, if the original tensor is

[[[A],
  [B], 
  [C]]]

(batch_size = 1, S = 3, A, B, C are sized 5)

and the list is [[0, 2]]
The result will be:

[[[A], 
  [C]]]

and that it’ll be possible to backpropagate through this tensor.

I heard about the gather function but i’m not sure if it fits here and how to use it in this case.

Thank you!

2

Hi, can you specify what you mean by “keeping the gradient of the original tensor”?

If it means that you just want to propagate the gradient incoming to:

[[[A], 
  [C]]]

we can note:

[[[grad(A)], 
  [grad(C)]]]

to the input of the operation in the most straight forward way you would obtain:

[[[grad(A)], 
  [None],
  [grad(C)]]]

Which is what would torch.gather do.

In case you want something else, please specify but I assume you will have to extend torch.autograd as in:
https://pytorch.org/docs/stable/notes/extending.html

Best,

Samuel

3

Hi and thank you for the reply.
I meant that i want to know the operation needed to get the tensor with the appropriate values from the original tensor & the list (as i explained), in such way that it’ll be possible to backpropagate through the result. (This computation is done while computing the loss so it has to be differential).
Thank you!

4

Ok in that case, torch.gather will do the job.

I think you could also do:

results = [[[A],
  [B], 
  [C]]]
l = [0, 2]
updatd_results = results[:, l, :]