Scatter scores in an adjacent matrix

autograd
1

I have one scores tensor with gradients attached to it and I have a corresponding lengths tensor where sum(lengths) = len(scores)

I would like to scatter these scores in a matrix in a diagonal way such that

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]

while maintaining the gradients so that it can flow back to the original scores tensor during backpropagation

image
image1322×1480 9.73 KB

Will appreciate any hints. Thanks

2

I figured a way but not sure how to vectorize this

            # 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
3

This should work, if I understand the use case correctly:

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)
> 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.]])
4

Yes repeat_interleave was the key. Thank you very much.