If the components are disjoint and you get the numbers in a single integer tensor, you could first use scatter_max from the PyTorch scatter library and then indexing with the mask to get the values back into the index.
There are faster ways to do this, but at least the one I implemented isn’t open source.
Best regards
Thomas