Solved by Ajinkya_Ambatwar in post #10 
              
              
                Thank you so much @ptrblck ! 
The segment_csr function from torch_scatter is just exactly what I wanted. 
from torch_tensor import segment_csr
x = torch.tensor([[6, 0, 3],
       [0, 7, 0],
       [4, 6, 8],
       [1, 6, 7],
       [2, 6, 1],
       [8, 5, 5],
       [9, 8, 3],
       [4, 9, 1],
 …