Modify array with list of indices

sachinruk: <p>Suppose I have a list of indices and wish to modify an existing array with this list. Currently the only way I can do this is by using a for loop as follows. Just wondering if there is a faster/ efficient way.</p>
<pre><code>torch.manual_seed(0)
a = torch.randn(5,3)
idx = torch.Tensor([[1,2], [3,2]]).to(torch.long)
for i,j in idx:
    a[i,j] = 1
</code></pre>
<p>I initially assumed that <code>gather</code> or <code>index_select</code> would go some way in answering this question, but looking at <a href="https://pytorch.org/docs/stable/torch.html#torch.gather" rel="nofollow noopener">documentation</a> this doesn’t seem to be the answer.</p>
<p>In my particular case, a is a 5 dimensional vector and idx is a Nx5 vector. So the output (after subscripting with something like <code>a[idx]</code>) I’d expect is a <code>(N,)</code> shaped vector.</p>
<p>This (unanswered) question is similar: <a href="https://discuss.pytorch.org/t/more-efficient-way-of-indexing-to-avoid-loop/27477">More efficient way of indexing to avoid loop</a></p>
<p>(this is a crosspost from <a href="https://stackoverflow.com/questions/52922401/pytorch-modify-array-with-list-of-indices" rel="nofollow noopener">stackoverflow</a>.)</p>

ptrblck: <p>This should work:</p>
<pre><code class="lang-python">a[idx[:, 0], idx[:, 1]] = 1.
</code></pre>

Xinyao_Wang: <p><a class="mention" href="/u/sachinruk">@sachinruk</a> Thanks for mentioning, and I would like to know how to extend this to higher dimension.<br>
If we have a 4-D tensor of [batch, channel, height, width], and a set of [x, y] coordinates with shape [batch, num_points, 2] , how to select this 4-D tensor without loop?</p>
<p>Here is my implementation with loop:</p>
<p>Suppose the 4-D feature map [10, 256, 64, 64] in dimension, and coordinates is [10, 68, 2] in dimension(for each batch, there are 68 points that we wanted to select from feature map)</p>
<pre><code class="lang-auto">        coord_features = torch.zeros(10, 68, 256)
        feature_map = feature_map.transpose(1,2).transpose(2,3) #reshape to [10, 64, 64, 256]
        for i in xrange(coords.shape[0]): #loop through each sample in batch
            for j in xrange(coords.shape[1]): #loop through each points
                #select coordinate on feature map
                coord_features[i][j] = feature_map[i][coords[i][j][1].type(torch.int64)][coords[i][j][0].type(torch.int64)]

</code></pre>

sachinruk: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>’s answer does work, but turns out this was the answer I was looking for as shown on SO:<br>
<code>a[idx.t().chunk(chunks=2,dim=0)] = 1</code></p>

pbloem: <p>To do this with a batch dimension, you can just create index tensors for the other dimensions as needed.</p>
<p>For a simple example, let’s say we have a tensor <code>data</code> of grayscale images with size <code>b, h, w</code> and a set of random pixels coordinates <code>coords</code> with size <code>b, 2</code> of pixel coordinates (one pixel per image in the batch) that we want to zero out. We can do this with:</p>
<pre><code class="lang-python">bselect = torch.arange(data.size(0), dtype=torch.long)
data[bselect, coords[:, 0], coords[:, 1]] = 0
</code></pre>
<p>Your example is a little more complicated, because you have multiple points per image, but you can just do something like <code>bselect = torch.range(batch, dtype=torch.long)[:, None].expand(batch, num_points).view(-1)</code>.</p>
<p>For the channel dimension, I think you can just use a single integer 0, but you may need to create a vector of zeros of the same length as <code>bselect</code>.</p>