How to efficiently perform averaging across predefined groups in a tensor

I am wondering is there a efficient way to perform group averaging in a tensor.
Take 1-D vector as a sample:
vector: [1, 2, 3, 4, 5, 6]
group_label:[1, 2, 1, 2, 3, 3]
results:[2, 3, 5.5 ]
The ideas is that values in a tensor with same group label will be averaged and produce a new tensor.

1 Like

You could use scatter_add and a small hack to get the unique counts of your indices:

x = torch.arange(1, 7, dtype=torch.float)
idx = torch.tensor([0, 1, 0, 1, 2, 2])
idx_unique = idx.unique(sorted=True)
idx_unique_count = torch.stack([(idx==idx_u).sum() for idx_u in idx_unique])
res = torch.zeros(len(idx_unique)).scatter_add(0, idx, x)
res /= idx_unique_count.float()

I think it’s time to add a return_counts option to torch.unqiue.

1 Like

Thanks for the help.

Thank you very much Sir.

thank you for your question!