Fixing edge_mask handling for directed graphs in k_hop_subgraph

[ryoji-kubo]

Example use case:

from torch_geometric.utils import k_hop_subgraph
import torch

edge_index = torch.tensor([[1, 2, 3],
                           [0, 1, 1]])
# get the 2-hop neighbors of node 0 in the directed graph.
_, edge_index, _, edge_mask = k_hop_subgraph(0, 2, edge_index, relabel_nodes=False, directed=True)

This gives the following result:

Expected Outcome:

>>> edge_index
tensor([[1, 2, 3],
        [0, 1, 1]])

>>> edge_mask
tensor([True,  True,  True])

Actual Outcome:

>>> edge_index
tensor([[2, 3],
        [1, 1]])
>>> edge_mask
tensor([False,  True,  True])

This stems from the fact that the line torch.index_select(node_mask, 0, row, out=edge_mask) overwrites edge_mask, effectively only marking the edges used in the final hop as True.

To fix this, I have added preserved_edge_mask that will mark the edges used in each hop as True.

[rusty1s]

Thanks :) Also added a test.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 87.33%. Comparing base (f5c8293) to head (923f35c).
Report is 68 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #9756      +/-   ##
==========================================
- Coverage   88.13%   87.33%   -0.81%     
==========================================
  Files         482      482              
  Lines       31421    31427       +6     
==========================================
- Hits        27693    27446     -247     
- Misses       3728     3981     +253     

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