How to select specific elements of a tensor?

seyeeet: <p>I have tensor</p>
<pre><code class="lang-auto">y=
torch.arange(50).reshape(5,10)
tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44, 45, 46, 47, 48, 49]])
</code></pre>
<p>and now I want to select elements of <code>y</code> that their indexes appears in <code>s</code> such that each row of <code>s</code> is correspond the the row of <code>y</code> that I want to selects the elements from. for example if s is</p>
<pre><code class="lang-auto">s = 
torch.arange(10).reshape(5,2)
tensor([[0, 1],
        [2, 3],
        [4, 5],
        [6, 7],
        [8, 9]])
</code></pre>
<p>the output after slicing (y[s]) should be:</p>
<pre><code class="lang-auto">[[0,1],
[12,13],
[24,25],
[36,37],
[48,49]]
</code></pre>
<ul>
<li>how can I do it?</li>
<li>I tried y[s], but it is not working, also I know that <code>s</code> should have same number of columns</li>
<li>will the gradient be reserved after slicing?</li>
</ul>

seyeeet: <p>my naive approch would be this:</p>
<pre><code class="lang-auto">stack = []
for i in range(s.shape[0]):
    stack.append(y[:,s][:,i,:][i,:])
output_selec = torch.stack(stack)
output_selec
tensor([[ 0,  1],
        [12, 13],
        [24, 25],
        [36, 37],
        [48, 49]])
</code></pre>
<p>but 1) I am sure there should be a more pythonic/cleaner way (maybe <a class="mention" href="/u/ptrblck">@ptrblck</a> can give me a suggestion), and 2) I am not sure if it will cause a disconnection between the gradients from previous operations</p>

abhibha1807: <p>As far as I remember, slicing tensors does cause the tensor to lose its grad_fn.</p>

seyeeet: <p>I see, any suggestion on how to slice and keep the grad_fn?</p>

eqy: <p>Does <code>torch.gather</code> work for you?</p>
<pre><code class="lang-auto">import torch
y = torch.arange(50).reshape(5,10)
s = torch.arange(10).reshape(5,2)
print(torch.gather(y, 1, s))
</code></pre>
<pre><code class="lang-auto">tensor([[ 0,  1],
        [12, 13],
        [24, 25],
        [36, 37],
        [48, 49]])
</code></pre>

seyeeet: <p>Apparently, it does, thanks.<br>
<a class="mention" href="/u/eqy">@eqy</a> would <code>torch.gather</code> safe for keeping the gradient?<br>
I see in the <code>torch.gather</code> <a href="https://pytorch.org/docs/stable/generated/torch.gather.html" rel="noopener nofollow ugc">documentations</a> that it has something called <code>sparse_grad</code>, I did not quite understand what does it mean to have <code>gradient to be a sparse tensor.</code>…</p>

eqy: <p>We can directly test if backprop works with both slice and gather:</p>
<pre><code class="lang-auto">import torch

class SliceNet(torch.nn.Module):
    def __init__(self):
        super(SliceNet, self).__init__()
        self.linear1 = torch.nn.Linear(32, 64)
        self.sigmoid = torch.nn.Sigmoid()
        self.linear2 = torch.nn.Linear(32, 1, bias=True)

    def forward(self, x):
        x = self.linear1(x)
        x = self.sigmoid(x)
        x = x[:,:32]
        x = self.linear2(x)
        return self.sigmoid(x)

class GatherNet(torch.nn.Module):
    def __init__(self):
        super(GatherNet, self).__init__()
        self.linear1 = torch.nn.Linear(32, 64)
        self.sigmoid = torch.nn.Sigmoid()
        self.linear2 = torch.nn.Linear(32, 1, bias=True)
        batch_size = 2
        self.indices = torch.arange(0, 64, 2).repeat(batch_size, 1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.sigmoid(x)
        x = torch.gather(x, 1, self.indices)
        x = self.linear2(x)
        return self.sigmoid(x)

def test_gather_gradient():
    print("testing gathernet, loss should go down")
    data = torch.randn(2, 32)
    label = torch.randn(2, 1).round()
    gathernet = GatherNet()
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(gathernet.parameters(), 1e-4)
    for i in range(10000):
        out = gathernet(data)
        loss = criterion(out, label)
        if i % 1000 == 0:
            print(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def test_slice_gradient():
    print("testing slicenet, loss should go down")
    data = torch.randn(2, 32)
    label = torch.randn(2, 1).round()
    slicenet = SliceNet()
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(slicenet.parameters(), 1e-4)
    for i in range(10000):
        out = slicenet(data)
        loss = criterion(out, label)
        if i % 1000 == 0:
            print(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

test_gather_gradient()
test_slice_gradient()
</code></pre>
<pre><code class="lang-auto">testing gathernet, loss should go down
0.16012126207351685
0.13146084547042847
0.10953329503536224
0.09261573851108551
0.07940082252025604
0.06893321871757507
0.06052353233098984
0.053674302995204926
0.04802417382597923
0.043307825922966
testing slicenet, loss should go down
1.1417878866195679
0.950451672077179
0.828807532787323
0.7515791654586792
0.7005079984664917
0.6650912761688232
0.6394497156143188
0.6201967597007751
0.60529625415802
0.593470573425293
</code></pre>
<p>In general, I don’t think most indexing operations should be problematic for backprop. For example, could max pooling work if indexing was a problem?</p>

ptrblck: <aside class="quote no-group" data-username="abhibha1807" data-post="3" data-topic="124375" 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/abhibha1807/48/38448_2.png" class="avatar"> abhibha1807:</div>
<blockquote>
<p>As far as I remember, slicing tensors does cause the tensor to lose its grad_fn.</p>
</blockquote>
</aside>
<p>Minor correction: slicing doesn’t detach the output:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn(4, 4, requires_grad=True)
y = x[:3, :2]
y.mean().backward()
print(x.grad)
&gt; tensor([[0.1667, 0.1667, 0.0000, 0.0000],
          [0.1667, 0.1667, 0.0000, 0.0000],
          [0.1667, 0.1667, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000]])
</code></pre>