title: Pytorch softmax along different masks without for loop
body: <p>Say I have a vector <code>a</code> , with an index vector <code>b</code> of the same length. The indexs are in range 0~N-1, corresponding to N groups. How can I do softmax for every group without for loop?</p>
  
  <p>I'm doing some sort of attention operation here. The numbers for every group are not the same, so I can't reshape <code>a</code> to a matrix and use the <code>dim</code> in standard <code>Softmax()</code> API.</p>
  
  <p>Toy example:</p>
  
  <pre><code>a = torch.rand(10)
  a: tensor([0.3376, 0.0557, 0.3016, 0.5550, 0.5814, 0.1306, 0.2697, 0.9989, 0.4917,
          0.6306])
  b = torch.randint(0,3,(1,10), dtype=torch.int64)
  b: tensor([[1, 2, 0, 2, 2, 0, 1, 1, 1, 1]])
  </code></pre>
  
  <p>I want to do softmax like</p>
  
  <pre><code>for index in range(3):
      softmax(a[b == index])
  </code></pre>
  
  <p>but without the for loop to save time.</p>