Switching Indices of Tensor

Jack_Rolph1: <p>I have a tensor, <code>A</code>, of shape <code>B</code> x <code>N</code> x  2. For instance:</p>
<pre><code class="lang-auto">B = 4
N = 5
A = torch.randn(B, N, 2) #(B x N x 2) Tensor
</code></pre>
<p>I then have a list of indices, that ought to sometimes ‘swap’ the positions of the last column of the tensor:</p>
<p><code>I = torch.tensor([[0,1], [1,0], [1,0], [0,1]])  #(B x 2) Tensor</code></p>
<p>The <code>[0,1]</code> indices ought to do nothing. The <code>[1,0]</code> indices ought to switch the last two columns.</p>
<p>How do I index <code>A</code> with <code>I</code> to perform this operation in Pytorch?</p>
<p><strong>Example:</strong></p>
<p>Suppose I have</p>
<pre><code class="lang-auto">print(A[0, :, :])

tensor([
        [1, 2],
        [3, 4],
        [5, 6],
        [7, 8],
        [9, 10]]])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">print(A[1, :, :])

tensor([
        [1, 2],
        [3, 4],
        [5, 6],
        [7, 8],
        [9, 10]]])
</code></pre>
<p>Then <code>foo(A, I)</code>, the indexing operation, should result in nothing happening to the first batch element (<code>[0,1]</code>), and a switch for the second batch element (<code>[1,0]</code>),:</p>
<pre><code class="lang-auto">print(foo(A, I)[0, :, :]))

tensor([
        [1, 2],
        [3, 4],
        [5, 6],
        [7, 8],
        [9, 10]]])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">print(foo(A, I)[1, :, :]))

tensor([
        [2, 1],
        [4, 3],
        [6, 5],
        [8, 7],
        [10, 9]]])
</code></pre>

KFrank: <p>Hi Jack!</p>
<aside class="quote no-group" data-username="Jack_Rolph1" data-post="1" data-topic="149650" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jack_rolph1/48/47137_2.png" class="avatar"> Jack_Rolph1:</div>
<blockquote>
<pre><code class="lang-auto">print(foo(A, I)[0, :, :]))

tensor([
        [1, 2],
        [3, 4],
        [5, 6],
        [7, 8],
        [9, 10]]])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">print(foo(A, I)[1, :, :]))

tensor([
        [2, 1],
        [4, 3],
        [6, 5],
        [8, 7],
        [10, 9]]])
</code></pre>
</blockquote>
</aside>
<p>You can use <code>torch.gather()</code> (after adjusting the shape of <code>I</code>):</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt; B = 4
&gt;&gt;&gt; N = 5
&gt;&gt;&gt; A = (torch.arange (N * 2) + 1).reshape (1, N, 2).expand (B, N, 2)
&gt;&gt;&gt; A
tensor([[[ 1,  2],
         [ 3,  4],
         [ 5,  6],
         [ 7,  8],
         [ 9, 10]],

        [[ 1,  2],
         [ 3,  4],
         [ 5,  6],
         [ 7,  8],
         [ 9, 10]],

        [[ 1,  2],
         [ 3,  4],
         [ 5,  6],
         [ 7,  8],
         [ 9, 10]],

        [[ 1,  2],
         [ 3,  4],
         [ 5,  6],
         [ 7,  8],
         [ 9, 10]]])
&gt;&gt;&gt; I = torch.tensor ([[0, 1], [1, 0], [1, 0], [0, 1]])
&gt;&gt;&gt; torch.gather (A, 2, I.unsqueeze (1).expand (4, 5, 2))
tensor([[[ 1,  2],
         [ 3,  4],
         [ 5,  6],
         [ 7,  8],
         [ 9, 10]],

        [[ 2,  1],
         [ 4,  3],
         [ 6,  5],
         [ 8,  7],
         [10,  9]],

        [[ 2,  1],
         [ 4,  3],
         [ 6,  5],
         [ 8,  7],
         [10,  9]],

        [[ 1,  2],
         [ 3,  4],
         [ 5,  6],
         [ 7,  8],
         [ 9, 10]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>