Is there any easy way to compute a ``head-mask-specified" pooling computation?

username: ryancc
cooked: <p>For example, I have a head-mask <code>h = [1, 1, 0, 0, 1]</code>, where <code>1</code> denotes the head of a span.<br>
And, I also have a tense <code>a = torch.tenser(5, 100)</code>. I want to implement a function, which should compute a new tensor <code>b</code> with a size of (3, 100), such that b[0] = a[0], b[1] = a[1] + a[2] + a[3] (as denoted in h), b[2] = a[4].<br>
Using <code>for</code> computation is easy, but is there any efficient computation? What if <code>h</code> and <code>a</code> are batched?</p>

username: Naruto-Sasuke
cooked: <p>Hi, ryancc<br>
Still confused about what you are talking about. Please clarify more clearly, maybe I can help.</p>

username: ryancc
cooked: <p>Hi, thank you very much.</p>
<p>Say I have a sentence consisting of two words: S = [“Definitely”, “not”], and what I want is to transfer S into an embedding matrix T with a size of (2, 100), where each row represents a word.</p>
<p>I want to adopt BERT embeddings. But in BERT, each word is represented as a sub-word unit. This means that S will be represented as [“Def”, “#<span class="hashtag">#in</span>”, “#<span class="hashtag">#ite</span>”, “#<span class="hashtag">#ly</span>”, “not”] ( “Definitely” is tokenized as “Def”, “#<span class="hashtag">#in</span>”, “#<span class="hashtag">#ite</span>”, “#<span class="hashtag">#ly</span>”). BERT will output an embedding matrix H with a size of (5, 100) :(.</p>
<p>My goal is to merge some rows of H according to the sub-word units.<br>
For example, for “Definitely”, I should merge the embeddings of [“Def”, “#<span class="hashtag">#in</span>”, “#<span class="hashtag">#ite</span>”, “#<span class="hashtag">#ly</span>”] to get its presentation.</p>
<p>In my current method, I use a head mask vector h = [1, 0, 0, 0, 1] to record the “head” of each word, where 1 indicates the head position:<br>
h = [<br>
1, -&gt; “Def”<br>
0, -&gt; “#<span class="hashtag">#in</span>”<br>
0, -&gt; “#<span class="hashtag">#ite</span>”<br>
0, -&gt; “#<span class="hashtag">#ly</span>”<br>
1 -&gt; “not”<br>
]<br>
So I should merge rows which have a head mask of 0 to that having a head mask of 1. I have to use the <code>for</code>  computation to enumerate each element in  <code>h</code> , which is slow and can not batchfy.</p>
<p>Could you give me some efficient method to do the above computation?</p>

username: Naruto-Sasuke
cooked: <p>I have found an efficient way to deal with it.</p>

username: Naruto-Sasuke
cooked: <pre><code class="lang-python">import torch

N = 6
C = 4

data = torch.randn(N, C)
head_mask = torch.tensor([1, 1, 0, 0, 1, 0])
head_index = (head_mask == 1).nonzero().view(-1) # 0, 1, 4
head_index_expand = torch.cat([head_index, torch.tensor([head_mask.numel()]) ] ) # 0,1, 4, 6
head_index_diff = (head_index_expand - torch.roll(head_index_expand, 1))[1:] # 1, 3, 2
splits = torch.split(data, head_index_diff.tolist())

print('ori data')
print(data)
print('after spliting')
print(splits)

# Each item contains a word, using combine all the elements in each item in the list
# you get the words you want.
</code></pre>

username: ryancc
cooked: <aside class="quote no-group" data-username="Naruto-Sasuke" data-post="6" data-topic="64854">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/ed8c4c/48.png" class="avatar"> Naruto-Sasuke:</div>
<blockquote>
<p>head_index_diff</p>
</blockquote>
</aside>
<p>Hi, Naruto-Sasuke:<br>
I have read your solution. It is exactly what I am looking for!<br>
Thank you very much!</p>