Add groupby op to support multi-version model inference

facebookresearch / rlmeta

RLMeta is a light-weight flexible framework for Distributed Reinforcement Learning Research.

MIT License

284 stars 28 forks source link

Add groupby op to support multi-version model inference #93

Closed xiaomengy closed 2 years ago

xiaomengy commented 2 years ago

This PR add a groupby op to group the indices of input tensor by values.

Example:

x = torch.tensor([0, 0, 1, 2, 1])
values, groups = rlmeta.ops.groupby(x)
# values: tensor([0, 1, 2])
# groups: (tensor([0, 1]), tensor([2, 4]), tensor([3]))