For example, I have a head-mask h = [1, 1, 0, 0, 1], where 1 denotes the head of a span.
And, I also have a tense a = torch.tenser(5, 100). I want to implement a function, which should compute a new tensor b 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].
Using for computation is easy, but is there any efficient computation? What if h and a are batched?