Slicing 3d tensor with multiple 2d tensors

username: MarkWallbang
cooked: <p>Hello,</p>
<p>I am implementing a Q-Learning algorithm for a routing problem at the moment, and I have come across a problem I don’t currently know how to solve. I have a 3d tensor that holds mutiple distance matrices of 2 dimensions (so <code> batchsize x coord1 x coord2</code>). Now I have a second tensor “B” that has batches of sequences of indices. Each row in each batch of “B” has indices as tuples, which refer to one cell in the distance matrix of the same batch. The dimension is hence <code>batchsize x rows x ntuples x 2</code> Now I want to select for each row in “B” (which consists of “n” tuples) the corresponding entries in the distance matrix of the same batch. I think torch.gather cannot help me there.<br>
A small example of batchsize 2 could look like this:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/5/0588f9e2bd5ae149e4885faab23f47e3b0db41f5.png" data-download-href="https://discuss.pytorch.org/uploads/default/0588f9e2bd5ae149e4885faab23f47e3b0db41f5" title="distmatrices"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/5/0588f9e2bd5ae149e4885faab23f47e3b0db41f5.png" alt="distmatrices" data-base62-sha1="MXR9z73IRaot5KmiQkP4l9I2MZ" width="690" height="409" data-dominant-color="F3F6F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">distmatrices</span><span class="informations">783×465 8.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Basically, each list of tuples should result in a flat list of the values from the corresponding distance matrix, so the output dimension is <code> batchsize x rows x ntuples</code> , since each index tuple resolves to one entry in the distance matrix. It is important that each batch of “B” should access the correct distance matrix, which is on the same dimension. I hope it was clear enough what I wanted to reach.</p>
<p>Any idea or help on how to solve this is well appreciated!<br>
Thank you in advance.</p>

username: ptrblck
cooked: <p>This code might work:</p>
<pre><code class="lang-python">idx = torch.tensor([[[[0, 1],
                      [1, 2],
                      [2, 3],
                      [3, 0]],
                     [[0, 1],
                      [1, 2],
                      [2, 3],
                      [3, 0]],
                     [[1, 0],
                      [0, 2],
                      [2, 3],
                      [3, 1]]],
                    [[[0, 0],
                      [1, 1],
                      [2, 2],
                      [3, 3]],
                     [[0, 0],
                      [1, 1],
                      [2, 2],
                      [3, 3]],
                     [[3, 3],
                      [2, 2],
                      [1, 1],
                      [0, 0]]]])

x = torch.randn(2, 4, 4)
x = x.unsqueeze(1).expand(-1, idx.size(1), -1, -1)
res = x[torch.arange(x.size(0))[:, None, None, None],
        torch.arange(x.size(1))[:, None, None],
        idx[:, :, :, 0:1], idx[:, :, :, 1:2]]
</code></pre>
<p>Could you verify it, as I was too lazy to type in your distance matrices? <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>

username: MarkWallbang
cooked: <p>Thank you for your fast response!<br>
The code seems to do exactly what I want. Is this easy to apply to different batch sizes?</p>

username: ptrblck
cooked: <p>Yes, different batch sizes should be directly usable, as I tried to avoid fixed batch sizes (or other dimensions).</p>
<p>Currently, of course <code>idx</code> is fixed as well as its last dimension, but I understand that that’s exactly the use case.</p>
<p>Let me know, if you encounter any issues using this code with other shapes.</p>