Vectorize for-loop - need to average slices of varying size

username: wetfire
cooked: <p>I am trying to average subword embeddings to form a word-level representation. Each word has a corresponding start and end index, indicating which subwords make up that word.</p>
<p><strong><code>sequence_output</code></strong> is a tensor of B * 384 * 768, where 384 is the max sequence length, and 768 is the number of features.</p>
<p><strong><code>all_token_mapping</code></strong> is a tensor of B * 384 * 2, which contains a start and end index. It is padded with [-1, -1].</p>
<p><strong><code>initial_reps</code></strong> is a tensor of num_nodes * 768, num_nodes is the sum of all the number of words (not subwords) in the different samples.</p>
<pre><code class="lang-auto">initial_reps = torch.empty((num_nodes, 768), dtype=torch.float32)
current_idx = 0
for i, feature_tokens_mapping in enumerate(all_token_mapping):
    for j, token_mapping in enumerate(feature_tokens_mapping):
        if token_mapping[0] == -1: # reached the end for this particular sequence
            break
        initial_reps[current_idx] = torch.mean(sequence_output[i][token_mapping[0]:token_mapping[-1] + 1], 0, keepdim=True)                                           
        current_idx += 1
</code></pre>
<p>My current code will create an empty tensor of length num_nodes, and a for loop will calculate the values at each index, by checking token_mapping[0] and token_mapping[1] for the correct slice of sequence_output to average.</p>
<p>Is there a way to vectorize this code?</p>
<p>In addition, I have a list that holds the number of words for each sample. i.e. the sum of all the elements in the list == num_nodes</p>
<p>Thank you.</p>

username: wetfire
cooked: <p>Not sure how to edit the original post so posting as comment instead.</p>
<p>Will use a simpler example.</p>
<p><strong><code>sequence_output</code></strong> is a tensor of B * 3 * 2, where 3 is the max sequence length, and 2 is the number of features.</p>
<p><strong><code>all_token_mapping</code></strong> is a tensor of B * 3 * 2, which contains a start and end index.</p>
<p><strong><code>initial_reps</code></strong> is a tensor of num_nodes * 2, num_nodes is the sum of all the number of words (not subwords) in the different samples.</p>
<pre><code class="lang-auto">sequence_output = torch.arange(2*3*2).float().reshape(2, 3, 2)
tensor([[[ 0.,  1.],
         [ 2.,  3.],
         [ 4.,  5.]],

        [[ 6.,  7.],
         [ 8.,  9.],
         [10., 11.]]])
</code></pre>
<pre><code class="lang-auto">all_token_mapping = torch.tensor([[[0,0],[1,2],[-1,-1]], [[0,2],[-1,-1],[-1,-1]]])
tensor([[[ 0,  0],
         [ 1,  2],
         [-1, -1]],

        [[ 0,  2],
         [-1, -1],
         [-1, -1]]])
</code></pre>
<pre><code class="lang-auto">num_nodes = 0
for sample in all_token_mapping:
  for mapping in sample:
    if mapping[0] != -1:
      num_nodes += 1
3
</code></pre>
<pre><code class="lang-auto">initial_reps = torch.empty((num_nodes, 2), dtype=torch.float32)
current_idx = 0
for i, feature_tokens_mapping in enumerate(all_token_mapping):
    for j, token_mapping in enumerate(feature_tokens_mapping):
        if token_mapping[0] == -1: # reached the end for this particular sequence
            break
        initial_reps[current_idx] = torch.mean(sequence_output[i][token_mapping[0]:token_mapping[-1] + 1], 0, keepdim=True)                                           
        current_idx += 1
initial_reps
tensor([[0., 1.],
        [3., 4.],
        [8., 9.]])
</code></pre>
<p>In the example above, initial_reps[0] will be the mean of sequence_output[0][0:1], initial_reps[1] will be the mean of sequence_output[0][1:3], and initial_reps[2] will be the mean of sequence_output[1][0:3].</p>
<p>My current code will create an empty tensor of length num_nodes, and a for loop will calculate the values at each index, by checking token_mapping[0] and token_mapping[1] for the correct slice of sequence_output to average.</p>
<p>Is there a way to vectorize this code?</p>
<p>In addition, I have a list that holds the number of words for each sample. i.e. the sum of all the elements in the list == num_nodes</p>

username: tom
cooked: <p>I think the encoding with “break” makes it a bit hard.<br>
If you encode sequence ids instead (can be done with comparison and cumsum if you don’t have it already), you can use index_add or some scatter_add implementation.<br>
Last I looked the public implementations for PyTorch were not terribly optimized but what can you do.</p>

username: wetfire
cooked: <p>Thanks to tom, I found out that scatter_add exists, and from there I found torch_scatter’s segment_coo<br>
Here’s my solution now:</p>
<pre><code>    initial_reps_list = []
    for i, sample_output in enumerate(sequence_output):
        token_mapping = all_token_mapping[i]
        token_mapping = token_mapping[token_mapping != -1]
        non_padded_outputs = sample_output[:num_bert_tokens[i]]
        initial_reps_list.append(torch_scatter.segment_coo(non_padded_outputs, token_mapping, reduce="mean"))
    initial_reps = torch.cat(initial_reps_list)
</code></pre>
<p>token_mapping is a list of indices in ascending order up to the max sequence length, padded with -1. I loop through the batch, for each sample, I get the token mapping, and only keep the non-negative indices.<br>
num_bert_tokens is a list that holds, for each sample, the number of tokens (no padding). I get the non-padded outputs, use segment_coo to reduce them according to the token_mapping, and append them all to a list.<br>
After the loop, I concatenate all the tensors in the list together.<br>
The method segment_coo reduces all values from the src tensor into out at the indices specified in the index tensor along the last dimension of index. More details can be found at: <a href="https://pytorch-scatter.readthedocs.io/en/latest/functions/segment_coo.html" class="inline-onebox" rel="noopener nofollow ugc">Segment COO — pytorch_scatter 2.0.6 documentation</a></p>
<p>Code runs much faster now <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>