nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: pyg-team/pytorch_geometric
Pull request #9807 (2024-11-26 14:12:05)

Check that custom edge types actually exist in `NumNeighbors` definition


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 341be665f..3867de8b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Dropped Python 3.8 support ([#9696](https://github.com/pyg-team/pytorch_geometric/pull/9606))
+- Added a check that confirms that custom edge types of `NumNeighbors` actually exist in the graph ([#9807](https://github.com/pyg-team/pytorch_geometric/pull/9807))
 
 ### Deprecated
 
diff --git a/torch_geometric/sampler/base.py b/torch_geometric/sampler/base.py
index d67ddd5af..1bd2e4346 100644
--- a/torch_geometric/sampler/base.py
+++ b/torch_geometric/sampler/base.py
@@ -425,6 +425,14 @@ class NumNeighbors:
             else:
                 assert False
 
+            # Confirm that `values` only hold valid edge types:
+            if isinstance(self.values, dict):
+                edge_types_str = {EdgeTypeStr(key) for key in edge_types}
+                invalid_edge_types = set(self.values.keys()) - edge_types_str
+                if len(invalid_edge_types) > 0:
+                    raise ValueError("Not all edge types specified in "
+                                     "'num_neighbors' exist in the graph")
+
             out = {}
             for edge_type in edge_types:
                 edge_type_str = EdgeTypeStr(edge_type)


Test patch:
diff --git a/test/sampler/test_sampler_base.py b/test/sampler/test_sampler_base.py
index dc8142176..41a7da255 100644
--- a/test/sampler/test_sampler_base.py
+++ b/test/sampler/test_sampler_base.py
@@ -49,6 +49,9 @@ def test_heterogeneous_num_neighbors_dict_and_default():
 
     num_neighbors = NumNeighbors({('A', 'B'): [25, 10]}, default=[-1, -1])
 
+    with pytest.raises(ValueError, match="Not all edge types"):
+        num_neighbors.get_values([('A', 'C'), ('B', 'A')])
+
     values = num_neighbors.get_values([('A', 'B'), ('B', 'A')])
     assert values == {('A', 'B'): [25, 10], ('B', 'A'): [-1, -1]}