Scatter scores in an adjacent matrix

username: Ahmed_Abdelaziz
cooked: <p>I have one scores tensor with gradients attached to it and I have a corresponding lengths tensor where <code>sum(lengths) = len(scores)</code></p>
<p>I would like to scatter these scores in a matrix in a diagonal way such that</p>
<pre data-code-wrap="python"><code class="lang-python">length_so_far = 0
for i in lengths:
    adj[i][length_so_far:length_so_far+lengths[i]] = scores[length_so_far:length_so_far+lengths[i]]
    length_so_far += lengths[i]
</code></pre>
<p>while maintaining the gradients so that it can flow back to the original scores tensor during backpropagation</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/f/5f3e2c8e1d1e51834197d72b0da0649991e3a6d4.png" data-download-href="https://discuss.pytorch.org/uploads/default/5f3e2c8e1d1e51834197d72b0da0649991e3a6d4" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f3e2c8e1d1e51834197d72b0da0649991e3a6d4_2_446x500.png" alt="image" data-base62-sha1="dAyz9pZbtwLXuWI5HXjrG8ObTOk" width="446" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f3e2c8e1d1e51834197d72b0da0649991e3a6d4_2_446x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f3e2c8e1d1e51834197d72b0da0649991e3a6d4_2_669x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f3e2c8e1d1e51834197d72b0da0649991e3a6d4_2_892x1000.png 2x" data-dominant-color="F9F9F9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1322×1480 9.73 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Will appreciate any hints. Thanks</p>

username: Ahmed_Abdelaziz
cooked: <p>I figured a way but not sure how to vectorize this</p>
<pre><code class="lang-python">            # Indexes to traverse rows
            z_index_1 = []

            # Indexes to traverse cols
             z_index_2 = []

            pointer = 0
            entity_pointer = 0
            
            # num_entities is a 1d tensor that is equivalent to lengths in figure
            for i in range(len(num_entities)):
                for j in range(num_entities[i]):
                    for k in range(num_entities[i]):
                        for l in range(8):
                            z_index_1.append(entity_pointer)
                            z_index_2.append(pointer)
                            pointer += 1
                    entity_pointer += 1

            z_scores_rearranged = torch.zeros(sum(num_entities), sum(num_entities * num_entities) * 8)

            z_scores_rearranged[(z_index_1, z_index_2)] = masked_flattened_scores
</code></pre>

username: ptrblck
cooked: <p>This should work, if I understand the use case correctly:</p>
<pre><code class="lang-python">lengths = torch.tensor([2, 3, 4, 5])
scores = torch.arange(1, lengths.sum()+1).float()

idx = torch.repeat_interleave(torch.arange(len(lengths)), lengths)

out = torch.zeros(len(lengths), len(scores))
out.scatter_(0, idx.unsqueeze(0), scores.unsqueeze(0))
print(out)
&gt; tensor([[ 1.,  2.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  3.,  4.,  5.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  6.,  7.,  8.,  9.,  0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10., 11., 12., 13., 14.]])
</code></pre>

username: Ahmed_Abdelaziz
cooked: <p>Yes <code>repeat_interleave</code> was the key. Thank you very much.</p>