Selecting elements in Tensor with differing amount of dimensions using another Tensor with indicies

LuggiStruggi: <p>So i have a tensor with a certain amount of dimensions <strong>A</strong> filled with values and a Tensor <strong>C</strong> with the same dimensions (but without the last)  with indices for <strong>A</strong>’s last dimension. I want to pick the values of <strong>A</strong> based on <strong>C</strong> such that a Tensor with the shape of <strong>C</strong> is returned.</p>
<p>Example:</p>
<p><strong>A</strong> has shape (7, 2, 4), <strong>C</strong> has shape (7, 2).</p>
<p><strong>A</strong>:</p>
<pre><code>    [[[-10.0491,  -4.9780,  -3.0346,  -1.0746],
     [  1.1812,   1.8627,   3.2540,   5.5354]],

    [[ -9.9464,  -4.9588,  -2.9927,  -0.8602],
     [  1.0148,   1.9898,   2.7656,   4.7714]],

    [[ -9.7778,  -5.0038,  -3.0378,  -1.0750],
     [  1.0365,   2.1866,   2.8971,   4.8669]],

    [[-10.1701,  -4.8115,  -3.0066,  -1.0485],
     [  0.7645,   1.8798,   3.0735,   4.7153]],

    [[ -9.8422,  -4.9419,  -3.1802,  -0.8486],
     [  0.8864,   1.7320,   3.3117,   4.8196]],

    [[-10.0794,  -4.9817,  -2.9976,  -0.9717],
     [  0.6711,   1.9173,   2.8586,   4.9084]],

    [[-10.0262,  -5.1335,  -2.9970,  -0.9397],
     [  1.2652,   1.9704,   3.0415,   5.8505]]]
</code></pre>
<p><strong>I</strong>:</p>
<pre><code>   [[0, 1],
    [0, 3],
    [0, 1],
    [0, 3],
    [1, 0],
    [2, 1],
    [2, 1]]
</code></pre>
<p>The wanted resulting tensor <strong>B</strong> with shape (7, 2) is:</p>
<pre><code>    [[-10.0491,   1.8627],
    [ -9.9464,   4.7714],
    [ -9.7778,   2.1866],
    [-10.1701,   4.7153],
    [ -4.9419,   0.8864],
    [ -2.9976,   1.9173],
    [ -2.9970,   1.9704]]
</code></pre>
<p>I want this to be possible for <strong>A</strong> being n-dimensional and <strong>C</strong> being (n-1)-dimensional.<br>
I had this solution for n=3:</p>
<p>b = a[torch.arange(c.shape[0]).unsqueeze(1), torch.arange(c.shape[1]), c]</p>
<p>This will obviously cause problems for a differnt amount of dimensions.<br>
Thank you for help!</p>

ptrblck: <p><code>gather</code> should work:</p>
<pre><code class="lang-python">A = torch.tensor([[[-10.0491,  -4.9780,  -3.0346,  -1.0746],
                   [  1.1812,   1.8627,   3.2540,   5.5354]],
             
                  [[ -9.9464,  -4.9588,  -2.9927,  -0.8602],
                   [  1.0148,   1.9898,   2.7656,   4.7714]],
            
                  [[ -9.7778,  -5.0038,  -3.0378,  -1.0750],
                   [  1.0365,   2.1866,   2.8971,   4.8669]],
            
                  [[-10.1701,  -4.8115,  -3.0066,  -1.0485],
                   [  0.7645,   1.8798,   3.0735,   4.7153]],
            
                  [[ -9.8422,  -4.9419,  -3.1802,  -0.8486],
                   [  0.8864,   1.7320,   3.3117,   4.8196]],
            
                  [[-10.0794,  -4.9817,  -2.9976,  -0.9717],
                   [  0.6711,   1.9173,   2.8586,   4.9084]],
            
                  [[-10.0262,  -5.1335,  -2.9970,  -0.9397],
                   [  1.2652,   1.9704,   3.0415,   5.8505]]])

C = torch.tensor([[0, 1],
                  [0, 3],
                  [0, 1],
                  [0, 3],
                  [1, 0],
                  [2, 1],
                  [2, 1]])

ref = torch.tensor([[-10.0491,   1.8627],
                    [ -9.9464,   4.7714],
                    [ -9.7778,   2.1866],
                    [-10.1701,   4.7153],
                    [ -4.9419,   0.8864],
                    [ -2.9976,   1.9173],
                    [ -2.9970,   1.9704]])

res = A.gather(2, C.unsqueeze(2)).squeeze(2)

print((res == ref).all())
# tensor(True)
</code></pre>

LuggiStruggi: <p>Perfect this is what i was looking for.  <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=12" title=":blush:" class="emoji" alt=":blush:" loading="lazy" width="20" height="20"> For other than three dimensions:</p>
<pre><code class="lang-auto">A.gather(-1, C.unsqueeze(-1)).squeeze(-1)
</code></pre>