Fill value to matrix based on index

username: KelleyYin
cooked: <p>I have a index matrix and a matrix filled by zero .<br>
For example :</p>
<pre><code class="lang-auto">index = Tensor([[ 0,  1],
                [ 1,  2],
                [ 4,  1]])
a = Tensor([[ 0.,  0.,  0.,  0.,  0.,  0.],
            [ 0.,  0.,  0.,  0.,  0.,  0.],
            [ 0.,  0.,  0.,  0.,  0.,  0.]])
</code></pre>
<p>Therefore, I would like to fill 1 to the zero matrix based on the index matrix .<br>
Like this :</p>
<pre><code class="lang-auto">a = Tensor([[ 1,  1,  0,  0.,  0.,  0.],
            [ 0.,  1,  1,  0.,  0.,  0.],
            [ 0.,  1,  0.,  0.,  1,  0.]])
</code></pre>
<p>Is there any easy way ?</p>

username: yunusemre
cooked: <p>You should use <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.index_fill_" rel="nofollow noopener">index_fill_()</a> (doc link).</p>
<pre><code class="lang-auto">for i, ind in enumerate(index):
    a[i].index_fill_(0, ind, 1)
</code></pre>

username: ptrblck
cooked: <p>Here is another approach without a for loop:</p>
<pre><code class="lang-python">a[torch.arange(a.size(0)).unsqueeze(1), index] = 1.
</code></pre>

username: addisonklinke
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> if <code>index</code> does not have the same number of rows as <code>a</code>, I think a more generic approach is needed. For instance</p>
<pre><code class="lang-auto">index = torch.tensor([
    [0, 1],
    [1, 2]])
a = torch.zeros(3, 6)
a[torch.arange(a.size(0)).unsqueeze(1), index] = 1  # IndexError
a[index[:, 0], index[:, 1]] = 1                     # Fixed
</code></pre>
<p>Maybe someone knows a more compact way to write the second approach, or one that extends easily to more than 2 dimensions? I have always been puzzled why <code>a[index] = 1</code> doesn’t work (rather than selecting the multi-dimensional indices, it simply makes all elements of <code>a</code> equal to 1)</p>

username: ptrblck
cooked: <p>Your operation won’t yield the same results.<br>
The original code snippet gives:</p>
<pre><code class="lang-python">index = torch.tensor([
    [0, 1],
    [1, 2],
    [4, 1]])
a = torch.zeros(3, 6)
a[torch.arange(a.size(0)).unsqueeze(1), index] = 1

tensor([[1., 1., 0., 0., 0., 0.],
        [0., 1., 1., 0., 0., 0.],
        [0., 1., 0., 0., 1., 0.]])
</code></pre>
<p>while your approach will fail with an <code>IndexError</code> as it’ll try to index <code>dim0</code> with <code>4</code> so they are not compatible.</p>
<pre><code class="lang-python">index = torch.tensor([
    [0, 1],
    [1, 2],
    [4, 1]])
a = torch.zeros(3, 6)
a[index[:, 0], index[:, 1]] = 1

# &gt; IndexError: index 4 is out of bounds for dimension 0 with size 3
</code></pre>

username: addisonklinke
cooked: <p>I had redefined <code>index</code> so that <code>[4, 1]</code> was not included. My intended point was that rather than assuming each row of <code>a</code> will have 2 elements equal to 1, there is a more general case where any subset of the rows and elements could be chosen. For that setup, I found it more convenient to treat <code>index</code> as a list of <code>(x, y)</code> coordinates i.e.</p>
<pre><code class="lang-auto">index = torch.tensor([
    [0, 1],  # First row, second column                      
    [1, 2],  # Second row, third column
             # Nothing in the 3rd row
])
</code></pre>

username: ptrblck
cooked: <p>Sure, but your code snippet even with the missing <code>[4, 1]</code> row wouldn’t return the same result, so it’s not a replacement:</p>
<pre><code class="lang-python">index = torch.tensor([
    [0, 1],
    [1, 2]])
a = torch.zeros(3, 6)
a[torch.arange(index.size(0)).unsqueeze(1), index] = 1
print(a)
&gt; tensor([[1., 1., 0., 0., 0., 0.],
          [0., 1., 1., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0.]])

b = torch.zeros(3, 6)
b[index[:, 0], index[:, 1]] = 1
print(b)
&gt; tensor([[0., 1., 0., 0., 0., 0.],
          [0., 0., 1., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0.]])
</code></pre>