Additional option for `torch_geometric.utils.subgraph`

🚀 The feature, motivation and pitch

I’ve come across dropout_node again and I found an unexpected behavior for this function. For an example provided in the doc-string:

        >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],
        ...                            [1, 0, 2, 1, 3, 2]])
        >>> edge_index, edge_mask, node_mask = dropout_node(edge_index)
        >>> edge_index
        tensor([[0, 1],
                [1, 0]])
        >>> edge_mask
        tensor([ True,  True, False, False, False, False])
        >>> node_mask
        tensor([ True,  True, False, False])

Actually, we need to remove all edges that contained the dropped nodes [0, 1], but subgraph only drops edges that both ends are in the subset [0,1]: https://github.com/pyg-team/pytorch_geometric/blob/081a964638459d1374f2e90de8ccfc9f0ea1f043/torch_geometric/utils/subgraph.py#L98

So in this case it should be | rather than & in the above code? Should we add an additional option to support this?

Alternatives

Add an option for subgraph to support | rather than & in some cases. https://github.com/pyg-team/pytorch_geometric/blob/081a964638459d1374f2e90de8ccfc9f0ea1f043/torch_geometric/utils/subgraph.py#L98

Additional context

No response

[rusty1s]

Sorry for the confusion. Don’t we already do this?

edge_mask = node_mask[edge_index[0]] & node_mask[edge_index[1]] 

only maintains the edges for which both endpoints are in the non-dropped node set.

[EdisonLeeeee]

Oh you are right, my bad. Everything looks okay, sorry for the wasted time here.