Hi there,

bit of an odd question but I’m wondering if this is possible to do with pytorch out of the box.

Basically I’m trying to replicate an operation I can do quite easily using pandas groupby operations.

Let’s say I have a tensor that’s full of people’s ages in a particular class which is simply identified by an integer. So in the following example the class identifiers are in column 0: 99 and 55. Note they are of different sizes. The ages are in column 1.

```
ages_by_class = [[ 99, 24 ],
[ 99, 13 ],
[ 55, 33 ], #<--- ages not necessarily sorted in any order apriori
[ 55, 43 ],
[ 55, 36 ],
```

I’m trying to get the indexes or boolean mask or values corresponding to the topk ages *within each group*. So for the above, the boolean mask would look something like

```
mask_top2 := [[ 1 ],
[ 1 ],
[ 0 ],
[ 1 ],
[ 1 ],
```

In pandas the solution is simply to groupby() the sorted class id column, and then get the .head(k) but it’s not clear how I could translate this logic into pytorch. Any help would be greatly appreciated!

cheers