Repository: pyg-team/pytorch_geometric
Pull request #7714 (2023-07-09 03:33:13)

Raise error when collecting non-existing attributes in `HeteroData`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7714?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7714](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7714?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4408148) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/a4447449a7ab185a643dba878e8548c452a62a26?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (a444744) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #7714   +/-   ##
=======================================
  Coverage   91.58%   91.58%           
=======================================
  Files         452      452           
  Lines       25534    25534           
=======================================
  Hits        23385    23385           
  Misses       2149     2149           
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7714?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/explain/explanation.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7714?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2V4cGxhaW4vZXhwbGFuYXRpb24ucHk=) | `95.20% <ø> (-0.04%)` | :arrow_down: |
| [torch\_geometric/data/hetero\_data.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7714?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvaGV0ZXJvX2RhdGEucHk=) | `90.07% <100.00%> (+0.01%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f0b0f6bba..c7029a2e5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -73,6 +73,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Raise error when collecting non-existing attributes in `HeteroData` ([#7714](https://github.com/pyg-team/pytorch_geometric/pull/7714))
 - Renamed `dest` argument to `dst` in `utils.geodesic_distance` ([#7708](https://github.com/pyg-team/pytorch_geometric/pull/7708))
 - Changed `add_random_edge` to only add true negative edges ([#7654](https://github.com/pyg-team/pytorch_geometric/pull/7654))
 - Allowed the usage of `BasicGNN` models in `DeepGraphInfomax` ([#7648](https://github.com/pyg-team/pytorch_geometric/pull/7648))
diff --git a/benchmark/training/training_benchmark.py b/benchmark/training/training_benchmark.py
index e17fd0368..0abc5cdab 100644
--- a/benchmark/training/training_benchmark.py
+++ b/benchmark/training/training_benchmark.py
@@ -36,7 +36,7 @@ def train_homo(model, loader, optimizer, device, progress_bar=True, desc="",
     for batch in loader:
         optimizer.zero_grad()
         batch = batch.to(device)
-        if hasattr(batch, 'adj_t'):
+        if 'adj_t' in batch:
             edge_index = batch.adj_t
         else:
             edge_index = batch.edge_index
@@ -67,7 +67,7 @@ def train_hetero(model, loader, optimizer, device, progress_bar=True, desc="",
     for batch in loader:
         optimizer.zero_grad()
         batch = batch.to(device)
-        if len(batch.adj_t_dict) > 0:
+        if 'adj_t' in batch:
             edge_index_dict = batch.adj_t_dict
         else:
             edge_index_dict = batch.edge_index_dict
diff --git a/benchmark/utils/hetero_gat.py b/benchmark/utils/hetero_gat.py
index b1bda5c5e..d3483e4ec 100644
--- a/benchmark/utils/hetero_gat.py
+++ b/benchmark/utils/hetero_gat.py
@@ -22,7 +22,7 @@ class HeteroGAT(torch.nn.Module):
             loader = tqdm(loader, desc="Inference")
         for batch in loader:
             batch = batch.to(device)
-            if len(batch.adj_t_dict) > 0:
+            if 'adj_t' in batch:
                 self.model(batch.x_dict, batch.adj_t_dict)
             else:
                 self.model(batch.x_dict, batch.edge_index_dict)
@@ -35,7 +35,7 @@ class HeteroGAT(torch.nn.Module):
             loader = tqdm(loader, desc="Evaluate")
         for batch in loader:
             batch = batch.to(device)
-            if len(batch.adj_t_dict) > 0:
+            if 'adj_t' in batch:
                 out = self.model(batch.x_dict, batch.adj_t_dict)
             else:
                 out = self.model(batch.x_dict, batch.edge_index_dict)
diff --git a/benchmark/utils/hetero_sage.py b/benchmark/utils/hetero_sage.py
index 8b71b15cd..1445ef614 100644
--- a/benchmark/utils/hetero_sage.py
+++ b/benchmark/utils/hetero_sage.py
@@ -21,7 +21,7 @@ class HeteroGraphSAGE(torch.nn.Module):
             loader = tqdm(loader, desc="Inference")
         for batch in loader:
             batch = batch.to(device)
-            if len(batch.adj_t_dict) > 0:
+            if 'adj_t' in batch:
                 self.model(batch.x_dict, batch.adj_t_dict)
             else:
                 self.model(batch.x_dict, batch.edge_index_dict)
@@ -34,7 +34,7 @@ class HeteroGraphSAGE(torch.nn.Module):
             loader = tqdm(loader, desc="Evaluate")
         for batch in loader:
             batch = batch.to(device)
-            if len(batch.adj_t_dict) > 0:
+            if 'adj_t' in batch:
                 out = self.model(batch.x_dict, batch.adj_t_dict)
             else:
                 out = self.model(batch.x_dict, batch.edge_index_dict)
diff --git a/benchmark/utils/utils.py b/benchmark/utils/utils.py
index ecb5d16aa..5b2dd3d9f 100644
--- a/benchmark/utils/utils.py
+++ b/benchmark/utils/utils.py
@@ -168,7 +168,7 @@ def test(model, loader, device, hetero, progress_bar=True,
     if hetero:
         for batch in loader:
             batch = batch.to(device)
-            if len(batch.adj_t_dict) > 0:
+            if 'adj_t' in batch:
                 edge_index_dict = batch.adj_t_dict
             else:
                 edge_index_dict = batch.edge_index_dict
@@ -182,7 +182,7 @@ def test(model, loader, device, hetero, progress_bar=True,
     else:
         for batch in loader:
             batch = batch.to(device)
-            if hasattr(batch, 'adj_t'):
+            if 'adj_t' in batch:
                 edge_index = batch.adj_t
             else:
                 edge_index = batch.edge_index
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 16c4320f5..cfe58a3b2 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -502,7 +502,11 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         """
         return self.node_types, self.edge_types
 
-    def collect(self, key: str) -> Dict[NodeOrEdgeType, Any]:
+    def collect(
+        self,
+        key: str,
+        allow_empty: bool = False,
+    ) -> Dict[NodeOrEdgeType, Any]:
         r"""Collects the attribute :attr:`key` from all node and edge types.
 
         .. code-block:: python
@@ -517,12 +521,21 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         .. note::
 
             This is equivalent to writing :obj:`data.x_dict`.
+
+        Args:
+            key (str): The attribute to collect from all node and ege types.
+            allow_empty (bool, optional): If set to :obj:`True`, will not raise
+                an error in case the attribute does not exit in any node or
+                edge type. (default: :obj:`False`)
         """
         mapping = {}
         for subtype, store in chain(self._node_store_dict.items(),
                                     self._edge_store_dict.items()):
             if hasattr(store, key):
                 mapping[subtype] = getattr(store, key)
+        if not allow_empty and len(mapping) == 0:
+            raise KeyError(f"Tried to collect '{key}' but did not find any "
+                           f"occurrences of it in any node and/or edge type")
         return mapping
 
     def _check_type_name(self, name: str):
@@ -888,7 +901,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
             value = torch.cat(values, dim) if len(values) > 1 else values[0]
             data[key] = value
 
-        if len(self.edge_label_index_dict) > 0:
+        if 'edge_label_index' in self:
             edge_label_index_dict = self.edge_label_index_dict
             for edge_type, edge_label_index in edge_label_index_dict.items():
                 edge_label_index = edge_label_index.clone()
@@ -1093,7 +1106,7 @@ def to_homogeneous_edge_index(
     cumsum = 0
     edge_indices: List[Tensor] = []
     edge_slices: Dict[EdgeType, Tuple[int, int]] = {}
-    for edge_type, edge_index in data.edge_index_dict.items():
+    for edge_type, edge_index in data.collect('edge_index', True).items():
         edge_index = offset_edge_index(node_slices, edge_type, edge_index)
         edge_indices.append(edge_index)
         edge_slices[edge_type] = (cumsum, cumsum + edge_index.size(1))
diff --git a/torch_geometric/explain/explanation.py b/torch_geometric/explain/explanation.py
index 57e4b98ed..049c618a1 100644
--- a/torch_geometric/explain/explanation.py
+++ b/torch_geometric/explain/explanation.py
@@ -271,11 +271,11 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
         return self._apply_masks(
             node_mask_dict={
                 key: mask.sum(dim=-1) > 0
-                for key, mask in self.node_mask_dict.items()
+                for key, mask in self.collect('node_mask', True).items()
             },
             edge_mask_dict={
                 key: mask > 0
-                for key, mask in self.edge_mask_dict.items()
+                for key, mask in self.collect('edge_mask', True).items()
             },
         )
 
@@ -285,11 +285,11 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
         return self._apply_masks(
             node_mask_dict={
                 key: mask.sum(dim=-1) == 0
-                for key, mask in self.node_mask_dict.items()
+                for key, mask in self.collect('node_mask', True).items()
             },
             edge_mask_dict={
                 key: mask == 0
-                for key, mask in self.edge_mask_dict.items()
+                for key, mask in self.collect('edge_mask', True).items()
             },
         )
 
@@ -328,10 +328,6 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
                 plots all features. (default: :obj:`None`)
         """
         node_mask_dict = self.node_mask_dict
-        if len(node_mask_dict) == 0:
-            raise ValueError(f"The attribute 'node_mask' is not available "
-                             f"in '{self.__class__.__name__}' "
-                             f"(got {self.available_explanations})")
         for node_mask in node_mask_dict.values():
             if node_mask.dim() != 2 or node_mask.size(1) <= 1:
                 raise ValueError(f"Cannot compute feature importance for "


Test patch:
diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py
index 75a37ee77..bea01b6ad 100644
--- a/test/data/test_hetero_data.py
+++ b/test/data/test_hetero_data.py
@@ -106,6 +106,8 @@ def test_hetero_data_to_from_dict():
 
 def test_hetero_data_functions():
     data = HeteroData()
+    with pytest.raises(KeyError, match="did not find any occurrences of it"):
+        data.collect('x')
     data['paper'].x = x_paper
     data['author'].x = x_author
     data['paper', 'paper'].edge_index = edge_index_paper_paper
diff --git a/test/explain/test_hetero_explanation.py b/test/explain/test_hetero_explanation.py
index f492407ca..ebbb3f098 100644
--- a/test/explain/test_hetero_explanation.py
+++ b/test/explain/test_hetero_explanation.py
@@ -136,7 +136,7 @@ def test_visualize_feature_importance(
     path = osp.join(tmp_path, 'feature_importance.png')
 
     if node_mask_type is None:
-        with pytest.raises(ValueError, match="node_mask' is not"):
+        with pytest.raises(KeyError, match="Tried to collect 'node_mask'"):
             explanation.visualize_feature_importance(path, top_k=top_k)
     else:
         explanation.visualize_feature_importance(path, top_k=top_k)