fixing edge mask 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.