Local maximums of sub-tensors by index tensor
I have a tensor x of shape (1,n), and another index tensor d of shape (1,k). I’m trying to find the maximums of k sub-tensors
x[0:d[0]], x[d[0]:d[1]], x[d[1]:d[2]], ..., x[d[-2]: d[-1]]

So the output is a tensor of shape (1,k) with k local maximums. I can implement a for loop, but that’s too slow. Can I do it in parallel in PyTorch (or Numpy)?

I found the answer thanks to user7138814. There is a SegmentCSR function in torch_scatter that does the job:
from torch_scatter import segment_csr

src = torch.randn(10, 6, 64)
indptr = torch.tensor([0, 2, 5, 6])
indptr = indptr.view(1, -1)  # Broadcasting in the first and last dim.

out = segment_csr(src, indptr, reduce=&quot;sum&quot;)

print(out.size())
torch.Size([10, 3, 64])

 output: torch.Size([10, 3, 64])
