Account for unsorted inputs when computing e_id in NeighborSampler (#7953)

Author: rusty1s

CHANGELOG.md

@@ -91,6 +91,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Fixed a bug in which `batch.e_id` was not correctly computed on unsorted graph inputs ([#7953](https://github.com/pyg-team/pytorch_geometric/pull/7953))
 - Fixed `from_networkx` conversion from `nx.stochastic_block_model` graphs ([#7941](https://github.com/pyg-team/pytorch_geometric/pull/7941))
 - Fixed the usage of `bias_initializer` in `HeteroLinear` ([#7923](https://github.com/pyg-team/pytorch_geometric/pull/7923))
 - Fixed broken links in `HGBDataset` ([#7907](https://github.com/pyg-team/pytorch_geometric/pull/7907))

test/explain/algorithm/test_graphmask_explainer.py

@@ -123,7 +123,7 @@ def test_graph_mask_explainer_binary_classification(
 
     explainer = Explainer(
         model=model,
-        algorithm=GraphMaskExplainer(2, epochs=5),
+        algorithm=GraphMaskExplainer(2, epochs=5, log=False),
         explanation_type=explanation_type,
         node_mask_type=node_mask_type,
         edge_mask_type=edge_mask_type,
@@ -171,7 +171,7 @@ def test_graph_mask_explainer_multiclass_classification(
 
     explainer = Explainer(
         model=model,
-        algorithm=GraphMaskExplainer(2, epochs=5),
+        algorithm=GraphMaskExplainer(2, epochs=5, log=False),
         explanation_type=explanation_type,
         node_mask_type=node_mask_type,
         edge_mask_type=edge_mask_type,
@@ -216,7 +216,7 @@ def test_graph_mask_explainer_regression(
 
     explainer = Explainer(
         model=model,
-        algorithm=GraphMaskExplainer(2, epochs=5),
+        algorithm=GraphMaskExplainer(2, epochs=5, log=False),
         explanation_type=explanation_type,
         node_mask_type=node_mask_type,
         edge_mask_type=edge_mask_type,

test/loader/test_link_neighbor_loader.py

@@ -570,3 +570,26 @@ def test_hetero_link_neighbor_loader_triplet(disjoint, temporal, amount):
             for i in range(batch_size):
                 assert (node_store.time[node_store.batch == i].max()
                         <= node_store.seed_time[i])
+
+
+@withPackage('pyg_lib')
+def test_link_neighbor_loader_mapping():
+    edge_index = torch.tensor([
+        [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 5],
+        [1, 2, 3, 4, 5, 8, 6, 7, 9, 10, 6, 11],
+    ])
+    data = Data(edge_index=edge_index, num_nodes=12)
+
+    loader = LinkNeighborLoader(
+        data,
+        edge_label_index=data.edge_index,
+        num_neighbors=[1],
+        batch_size=2,
+        shuffle=True,
+    )
+
+    for batch in loader:
+        assert torch.equal(
+            batch.n_id[batch.edge_index],
+            data.edge_index[:, batch.e_id],
+        )

test/loader/test_neighbor_loader.py

@@ -692,3 +692,25 @@ def test_hetero_neighbor_loader_sampled_info():
     for edge_type in batch.edge_types:
         assert (batch[edge_type].num_sampled_edges ==
                 expected_num_sampled_edges[edge_type])
+
+
+@withPackage('pyg_lib')
+def test_neighbor_loader_mapping():
+    edge_index = torch.tensor([
+        [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 5],
+        [1, 2, 3, 4, 5, 8, 6, 7, 9, 10, 6, 11],
+    ])
+    data = Data(edge_index=edge_index, num_nodes=12)
+
+    loader = NeighborLoader(
+        data,
+        num_neighbors=[1],
+        batch_size=2,
+        shuffle=True,
+    )
+
+    for batch in loader:
+        assert torch.equal(
+            batch.n_id[batch.edge_index],
+            data.edge_index[:, batch.e_id],
+        )

torch_geometric/loader/link_loader.py

@@ -224,7 +224,9 @@ def filter_fn(
             if 'n_id' not in data:
                 data.n_id = out.node
             if out.edge is not None and 'e_id' not in data:
-                data.e_id = out.edge
+                edge = out.edge.to(torch.long)
+                perm = self.link_sampler.edge_permutation
+                data.e_id = perm[out.edge] if perm is not None else out.edge
 
             data.batch = out.batch
             data.num_sampled_nodes = out.num_sampled_nodes
@@ -260,8 +262,10 @@ def filter_fn(
                     data[key].n_id = node
 
             for key, edge in (out.edge or {}).items():
-                if 'e_id' not in data[key]:
-                    data[key].e_id = edge
+                if edge is not None and 'e_id' not in data[key]:
+                    edge = edge.to(torch.long)
+                    perm = self.link_sampler.edge_permutation[key]
+                    data[key].e_id = perm[edge] if perm is not None else edge
 
             data.set_value_dict('batch', out.batch)
             data.set_value_dict('num_sampled_nodes', out.num_sampled_nodes)

torch_geometric/loader/link_neighbor_loader.py

@@ -63,6 +63,14 @@ class LinkNeighborLoader(LinkLoader):
     The rest of the functionality mirrors that of
     :class:`~torch_geometric.loader.NeighborLoader`, including support for
     heterogeneous graphs.
+    In particular, the data loader will add the following attributes to the
+    returned mini-batch:
+
+    * :obj:`n_id` The global node index for every sampled node
+    * :obj:`e_id` The global edge index for every sampled edge
+    * :obj:`input_id`: The global index of the :obj:`edge_label_index`
+    * :obj:`num_sampled_nodes`: The number of sampled nodes in each hop
+    * :obj:`num_sampled_edges`: The number of sampled edges in each hop
 
     .. note::
         Negative sampling is currently implemented in an approximate

torch_geometric/loader/neighbor_loader.py

@@ -101,6 +101,16 @@ class NeighborLoader(NodeLoader):
         sampled_data = next(iter(loader))
         print(sampled_data.n_id)  # Global node index of each node in batch.
 
+    In particular, the data loader will add the following attributes to the
+    returned mini-batch:
+
+    * :obj:`batch_size` The number of seed nodes (first nodes in the batch)
+    * :obj:`n_id` The global node index for every sampled node
+    * :obj:`e_id` The global edge index for every sampled edge
+    * :obj:`input_id`: The global index of the :obj:`input_nodes`
+    * :obj:`num_sampled_nodes`: The number of sampled nodes in each hop
+    * :obj:`num_sampled_edges`: The number of sampled edges in each hop
+
     Args:
         data (Any): A :class:`~torch_geometric.data.Data`,
             :class:`~torch_geometric.data.HeteroData`, or

torch_geometric/loader/node_loader.py

@@ -156,7 +156,9 @@ def filter_fn(
             if 'n_id' not in data:
                 data.n_id = out.node
             if out.edge is not None and 'e_id' not in data:
-                data.e_id = out.edge
+                edge = out.edge.to(torch.long)
+                perm = self.node_sampler.edge_permutation
+                data.e_id = perm[edge] if perm is not None else edge
 
             data.batch = out.batch
             data.num_sampled_nodes = out.num_sampled_nodes
@@ -180,8 +182,10 @@ def filter_fn(
                     data[key].n_id = node
 
             for key, edge in (out.edge or {}).items():
-                if 'e_id' not in data[key]:
-                    data[key].e_id = edge
+                if edge is not None and 'e_id' not in data[key]:
+                    edge = edge.to(torch.long)
+                    perm = self.node_sampler.edge_permutation[key]
+                    data[key].e_id = perm[edge] if perm is not None else edge
 
             data.set_value_dict('batch', out.batch)
             data.set_value_dict('num_sampled_nodes', out.num_sampled_nodes)