Add faster mask_select method

[toenshoff]

This PR addresses a minor performance issue with the utils.mask.mask_select method.

The previous implementation unsqueezed the mask to be broadcastable along all dimensions.
This broadcasted masking is significantly slower in torch than applying a 1-dimensional mask to the first dimension of a tensor.
The new implementation simply transposes the tensor such that the masked dimension comes first and later undoes this.

This mainly affects the methods for computing subgraphs which are needlessly slow on bigger graphs with many features.

Here is a code example to illustrate the difference:

import torch
from torch_geometric.data import Data
from timeit import timeit


N = 1000000
E = 10000000

data = Data(
    x=torch.randn((N, 200)),
    edge_index=torch.randint(0, N, (2, E)),
    edge_attr=torch.randn((E, 200), dtype=torch.float16)
)

subset_bool = torch.rand((N,)) < 0.5
time_bool = timeit(lambda: data.subgraph(subset_bool), number=1)
print(f'Time: {time_bool:.4f}s')

The prior version yields:

Time: 13.7826s

The updated version yields:

Time:  0.3541s

So more than a 30-fold speedup. We thought this may be helpful to accelerate workflows involving subgraph extraction.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (b1f8535) 88.80% compared to head (3a2d630) 88.41%.
Report is 1 commits behind head on master.

❗ Current head 3a2d630 differs from pull request most recent head 959e03b. Consider uploading reports for the commit 959e03b to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #8369      +/-   ##
==========================================
- Coverage   88.80%   88.41%   -0.39%     
==========================================
  Files         475      475              
  Lines       28841    28837       -4     
==========================================
- Hits        25611    25497     -114     
- Misses       3230     3340     +110     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[akihironitta]

LGTM!

[rusty1s]

Great catch, thank you!