Fixed bug in utils.bipartite_subgraph

[toenshoff]

The PR fixes a bug in the utils.bipartite_subgraph method, which crashes if src_subset contains Boolean values but dst_subset contains integers.

Consider the following example:

import torch
from torch_geometric.utils import bipartite_subgraph

subset = (torch.tensor([True, True, False]), torch.tensor([0, 3]))
edge_index = torch.tensor([[0, 1, 2], [0, 0, 3]])
sub_edge_index = bipartite_subgraph(subset, edge_index)

With torch_geometric==2.3.0 we obtain the following error:

Traceback (most recent call last):
  File "/home/jan/git/gerd/pygtest.py", line 8, in <module>
    sub_edge_index = bipartite_subgraph(subset, edge_index)
  File "/home/jan/miniconda3/envs/gerd/lib/python3.9/site-packages/torch_geometric/utils/subgraph.py", line 178, in bipartite_subgraph
    edge_mask = src_subset[edge_index[0]] & dst_subset[edge_index[1]]
IndexError: index 3 is out of bounds for dimension 0 with size 2

The method determines whether or not a cast of the masks is needed by only checking the dtype of the source mask, which causes this issue.

While this seems like a rare edge case, this issue does cause the HeteroData.subgraph method to crash when subset_dict uses Boolean masks and does not contain some ntypes as keys, as integer tensors are then added for the missing ntypes.

Since #6218, HeteroData.subgraph does retain ntypes that are absent from subset_dict. I have also updated the doc string of the method to reflect this change, since it was still describing the previous behavior.

[codecov]

Codecov Report

Merging #6654 (1d585e2) into master (6196ab3) will increase coverage by 0.00%.
The diff coverage is 100.00%.

❗ Current head 1d585e2 differs from pull request most recent head a1585e6. Consider uploading reports for the commit a1585e6 to get more accurate results

@@           Coverage Diff           @@
##           master    #6654   +/-   ##
=======================================
  Coverage   88.97%   88.97%           
=======================================
  Files         423      423           
  Lines       22957    22958    +1     
=======================================
+ Hits        20425    20427    +2     
+ Misses       2532     2531    -1     

Impacted Files	Coverage Δ
torch_geometric/data/hetero_data.py	90.81% <ø> (ø)
torch_geometric/utils/subgraph.py	100.00% <100.00%> (ø)
torch_geometric/data/storage.py	82.68% <0.00%> (+0.24%)	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[wsad1]

Thanks for fixing.