Hotfix: #4856 follow-up, also fix edge_types (#4857)

* fix

* fix test

Co-authored-by: rusty1s <[email protected]>

Author: ZenoTan

CHANGELOG.md

@@ -6,7 +6,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [2.0.5] - 2022-MM-DD
 ### Added
 - Added a `NeighborLoader` benchmark script ([#4815](https://github.com/pyg-team/pytorch_geometric/pull/4815))
-- Added support for `FeatureStore` and `GraphStore` in `NeighborLoader` ([#4817](https://github.com/pyg-team/pytorch_geometric/pull/4817), [#4851](https://github.com/pyg-team/pytorch_geometric/pull/4851), [#4854](https://github.com/pyg-team/pytorch_geometric/pull/4854))
+- Added support for `FeatureStore` and `GraphStore` in `NeighborLoader` ([#4817](https://github.com/pyg-team/pytorch_geometric/pull/4817), [#4851](https://github.com/pyg-team/pytorch_geometric/pull/4851), [#4854](https://github.com/pyg-team/pytorch_geometric/pull/4854), [#4856](https://github.com/pyg-team/pytorch_geometric/pull/4856), [#4857](https://github.com/pyg-team/pytorch_geometric/pull/4857))
 - Added a `normalize` parameter to `dense_diff_pool` ([#4847](https://github.com/pyg-team/pytorch_geometric/pull/4847))
 - Added `size=None` explanation to jittable `MessagePassing` modules in the documentation ([#4850](https://github.com/pyg-team/pytorch_geometric/pull/4850))
 - Added documentation to the `DataLoaderIterator` class ([#4838](https://github.com/pyg-team/pytorch_geometric/pull/4838))

test/loader/test_neighbor_loader.py

@@ -336,12 +336,16 @@ def test_custom_neighbor_loader(FeatureStore, GraphStore):
     for batch1, batch2 in zip(loader1, loader2):
         assert len(batch1) == len(batch2)
         assert batch1['paper'].batch_size == batch2['paper'].batch_size
-        assert torch.allclose(batch1['paper'].x, batch2['paper'].x)
-        assert torch.allclose(batch1['author'].x, batch2['author'].x)
-
-        assert torch.allclose(batch1['paper', 'to', 'paper'].edge_index,
-                              batch2['paper', 'to', 'paper'].edge_index)
-        assert torch.allclose(batch1['paper', 'to', 'author'].edge_index,
-                              batch2['paper', 'to', 'author'].edge_index)
-        assert torch.allclose(batch1['author', 'to', 'paper'].edge_index,
-                              batch2['author', 'to', 'paper'].edge_index)
+
+        # Mapped indices of neighbors may be differently sorted:
+        assert torch.allclose(batch1['paper'].x.sort()[0],
+                              batch2['paper'].x.sort()[0])
+        assert torch.allclose(batch1['author'].x.sort()[0],
+                              batch2['author'].x.sort()[0])
+
+        assert (batch1['paper', 'to', 'paper'].edge_index.size() == batch1[
+            'paper', 'to', 'paper'].edge_index.size())
+        assert (batch1['paper', 'to', 'author'].edge_index.size() == batch1[
+            'paper', 'to', 'author'].edge_index.size())
+        assert (batch1['author', 'to', 'paper'].edge_index.size() == batch1[
+            'author', 'to', 'paper'].edge_index.size())

torch_geometric/loader/neighbor_loader.py

@@ -110,7 +110,8 @@ def __init__(
 
             self.node_types = list(
                 set(node_attr.group_name for node_attr in node_attrs))
-            self.edge_types = [edge_attr.edge_type for edge_attr in edge_attrs]
+            self.edge_types = list(
+                set(edge_attr.edge_type for edge_attr in edge_attrs))
 
             # Set other required parameters:
             if isinstance(num_neighbors, (list, tuple)):