Repository: pyg-team/pytorch_geometric
Pull request #9939 (2025-01-14 03:46:28)

Introduce `LinkPredMetric._prepare`; allow for less predictions than `k`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7321dbb30..8691f4af5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed the `k_hop_subgraph()` method for directed graphs ([#9756](https://github.com/pyg-team/pytorch_geometric/pull/9756))
 - Fixed `utils.group_cat` concatenating dimension ([#9766](https://github.com/pyg-team/pytorch_geometric/pull/9766))
 - Fixed `WebQSDataset.process` raising exceptions ([#9665](https://github.com/pyg-team/pytorch_geometric/pull/9665))
+- Fixed `is_node_attr()` and `is_edge_attr()` errors when `cat_dim` is a tuple ([#9895](https://github.com/pyg-team/pytorch_geometric/issues/9895))
 
 ### Removed
 
diff --git a/examples/llm/molecule_gpt.py b/examples/llm/molecule_gpt.py
index 8f6c60240..6f11d8796 100644
--- a/examples/llm/molecule_gpt.py
+++ b/examples/llm/molecule_gpt.py
@@ -167,7 +167,7 @@ def train(
                 f'moleculegpt_epoch{best_epoch}_val_loss{best_val_loss:4f}_ckpt.pt'  # noqa: E501
             )
     torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
 
     print(f"Total Training Time: {time.time() - start_time:2f}s")
     # Test
diff --git a/readthedocs.yml b/readthedocs.yml
index 44ed0553c..9c60a9b18 100644
--- a/readthedocs.yml
+++ b/readthedocs.yml
@@ -1,5 +1,8 @@
 version: 2
 
+sphinx:
+   configuration: docs/source/conf.py
+
 build:
    os: ubuntu-22.04
    tools:
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index 07a52fdfc..e39bcebd6 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -806,6 +806,10 @@ class GlobalStorage(NodeStorage, EdgeStorage):
             return False
 
         cat_dim = self._parent().__cat_dim__(key, value, self)
+
+        if not isinstance(cat_dim, int):
+            return False
+
         num_nodes, num_edges = self.num_nodes, self.num_edges
 
         if value.shape[cat_dim] != num_nodes:
@@ -852,6 +856,10 @@ class GlobalStorage(NodeStorage, EdgeStorage):
             return False
 
         cat_dim = self._parent().__cat_dim__(key, value, self)
+
+        if not isinstance(cat_dim, int):
+            return False
+
         num_nodes, num_edges = self.num_nodes, self.num_edges
 
         if value.shape[cat_dim] != num_edges:
diff --git a/torch_geometric/metrics/link_pred.py b/torch_geometric/metrics/link_pred.py
index 8785cf487..a450bcd1a 100644
--- a/torch_geometric/metrics/link_pred.py
+++ b/torch_geometric/metrics/link_pred.py
@@ -43,34 +43,15 @@ class LinkPredMetric(BaseMetric):
             self.register_buffer('accum', torch.tensor(0.))
             self.register_buffer('total', torch.tensor(0))
 
-    def update(
-        self,
+    @staticmethod
+    def _prepare(
         pred_index_mat: Tensor,
         edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
-    ) -> None:
-        r"""Updates the state variables based on the current mini-batch
-        prediction.
-
-        :meth:`update` can be repeated multiple times to accumulate the results
-        of successive predictions, *e.g.*, inside a mini-batch training or
-        evaluation loop.
-
-        Args:
-            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
-                every example in the mini-batch with shape
-                :obj:`[batch_size, k]`.
-            edge_label_index (torch.Tensor): The ground-truth indices for every
-                example in the mini-batch, given in COO format of shape
-                :obj:`[2, num_ground_truth_indices]`.
-        """
-        if pred_index_mat.size(1) != self.k:
-            raise ValueError(f"Expected 'pred_index_mat' to hold {self.k} "
-                             f"many indices for every entry "
-                             f"(got {pred_index_mat.size(1)})")
-
-        # Compute a boolean matrix indicating if the k-th prediction is part of
-        # the ground-truth. We do this by flattening both prediction and
-        # target indices, and then determining overlaps via `torch.isin`.
+    ) -> Tuple[Tensor, Tensor]:
+        # Compute a boolean matrix indicating if the `k`-th prediction is part
+        # of the ground-truth, as well as the number of ground-truths for every
+        # example. We do this by flattening both prediction and ground-truth
+        # indices, and then determining overlaps via `torch.isin`.
         max_index = max(  # type: ignore
             pred_index_mat.max() if pred_index_mat.numel() > 0 else 0,
             edge_label_index[1].max()
@@ -88,7 +69,7 @@ class LinkPredMetric(BaseMetric):
         pred_isin_mat = torch.isin(flat_pred_index, flat_y_index)
         pred_isin_mat = pred_isin_mat.view(pred_index_mat.size())
 
-        # Compute the number of targets per example:
+        # Compute the number of ground-truths per example:
         y_count = scatter(
             torch.ones_like(edge_label_index[0]),
             edge_label_index[0],
@@ -97,11 +78,41 @@ class LinkPredMetric(BaseMetric):
             reduce='sum',
         )
 
-        metric = self._compute(pred_isin_mat, y_count)
+        return pred_isin_mat, y_count
 
+    def _update_from_prepared(
+        self,
+        pred_isin_mat: Tensor,
+        y_count: Tensor,
+    ) -> None:
+        metric = self._compute(pred_isin_mat[:, :self.k], y_count)
         self.accum += metric.sum()
         self.total += (y_count > 0).sum()
 
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+    ) -> None:
+        r"""Updates the state variables based on the current mini-batch
+        prediction.
+
+        :meth:`update` can be repeated multiple times to accumulate the results
+        of successive predictions, *e.g.*, inside a mini-batch training or
+        evaluation loop.
+
+        Args:
+            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
+                every example in the mini-batch with shape
+                :obj:`[batch_size, k]`.
+            edge_label_index (torch.Tensor): The ground-truth indices for every
+                example in the mini-batch, given in COO format of shape
+                :obj:`[2, num_ground_truth_indices]`.
+        """
+        pred_isin_mat, y_count = self._prepare(pred_index_mat,
+                                               edge_label_index)
+        self._update_from_prepared(pred_isin_mat, y_count)
+
     def compute(self) -> Tensor:
         r"""Computes the final metric value."""
         if self.total == 0:
@@ -182,8 +193,9 @@ class LinkPredMAP(LinkPredMetric):
     higher_is_better: bool = True
 
     def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        cum_precision = (torch.cumsum(pred_isin_mat, dim=1) /
-                         torch.arange(1, self.k + 1, device=y_count.device))
+        device = pred_isin_mat.device
+        arange = torch.arange(1, pred_isin_mat.size(1) + 1, device=device)
+        cum_precision = pred_isin_mat.cumsum(dim=1) / arange
         return ((cum_precision * pred_isin_mat).sum(dim=-1) /
                 y_count.clamp(min=1e-7, max=self.k))
 
@@ -210,7 +222,8 @@ class LinkPredNDCG(LinkPredMetric):
         self.register_buffer('idcg', cumsum(multiplier))
 
     def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        dcg = (pred_isin_mat * self.multiplier.view(1, -1)).sum(dim=-1)
+        multiplier = self.multiplier[:pred_isin_mat.size(1)].view(1, -1)
+        dcg = (pred_isin_mat * multiplier).sum(dim=-1)
         idcg = self.idcg[y_count.clamp(max=self.k)]
 
         out = dcg / idcg
@@ -228,8 +241,6 @@ class LinkPredMRR(LinkPredMetric):
     higher_is_better: bool = True
 
     def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        rank = pred_isin_mat.type(torch.uint8).argmax(dim=-1)
-        is_correct = pred_isin_mat.gather(1, rank.view(-1, 1)).view(-1)
-        reciprocals = 1.0 / (rank + 1)
-        reciprocals[~is_correct] = 0.0
-        return reciprocals
+        device = pred_isin_mat.device
+        arange = torch.arange(1, pred_isin_mat.size(1) + 1, device=device)
+        return (pred_isin_mat / arange).max(dim=-1)[0]


Test patch:
diff --git a/test/data/lightning/test_datamodule.py b/test/data/lightning/test_datamodule.py
index 8f66e0591..00f97dc3c 100644
--- a/test/data/lightning/test_datamodule.py
+++ b/test/data/lightning/test_datamodule.py
@@ -115,17 +115,18 @@ def test_lightning_dataset(get_dataset, strategy_type):
                                       num_workers=3, shuffle=True)
         assert 'shuffle' not in datamodule.kwargs
     old_x = train_dataset._data.x.clone()
-    new_datamodule_repr = has_package('pytorch_lightning>=2.5.0')
-    datamodule_repr = ('{Train dataloader: size=50}\n'
-                       '{Validation dataloader: size=30}\n'
-                       '{Test dataloader: size=10}\n'
-                       '{Predict dataloader: size=98}' if new_datamodule_repr
-                       else 'LightningDataset(train_dataset=MUTAG(50), '
-                       'val_dataset=MUTAG(30), '
-                       'test_dataset=MUTAG(10), '
-                       'pred_dataset=MUTAG(98), batch_size=5, '
-                       'num_workers=3, pin_memory=True, '
-                       'persistent_workers=True)')
+    if has_package('pytorch_lightning>=2.5.0'):
+        datamodule_repr = ('{Train dataloader: size=50}\n'
+                           '{Validation dataloader: size=30}\n'
+                           '{Test dataloader: size=10}\n'
+                           '{Predict dataloader: size=98}')
+    else:
+        datamodule_repr = ('LightningDataset(train_dataset=MUTAG(50), '
+                           'val_dataset=MUTAG(30), '
+                           'test_dataset=MUTAG(10), '
+                           'pred_dataset=MUTAG(98), batch_size=5, '
+                           'num_workers=3, pin_memory=True, '
+                           'persistent_workers=True)')
     assert str(datamodule) == datamodule_repr
 
     trainer.fit(model, datamodule)
@@ -141,14 +142,16 @@ def test_lightning_dataset(get_dataset, strategy_type):
                              log_every_n_steps=1)
 
         datamodule = LightningDataset(train_dataset, batch_size=5)
-        datamodule_repr = ('{Train dataloader: size=50}\n'
-                           '{Validation dataloader: None}\n'
-                           '{Test dataloader: None}\n{'
-                           'Predict dataloader: None}' if new_datamodule_repr
-                           else 'LightningDataset(train_dataset=MUTAG(50), '
-                           'batch_size=5, num_workers=0, '
-                           'pin_memory=True, '
-                           'persistent_workers=False)')
+        if has_package('pytorch_lightning>=2.5.0'):
+            datamodule_repr = ('{Train dataloader: size=50}\n'
+                               '{Validation dataloader: None}\n'
+                               '{Test dataloader: None}\n{'
+                               'Predict dataloader: None}')
+        else:
+            datamodule_repr = ('LightningDataset(train_dataset=MUTAG(50), '
+                               'batch_size=5, num_workers=0, '
+                               'pin_memory=True, '
+                               'persistent_workers=False)')
         assert str(datamodule) == datamodule_repr
 
         with expect_rank_zero_user_warning("defined a `validation_step`"):
@@ -244,18 +247,19 @@ def test_lightning_node_data(get_dataset, strategy_type, loader):
                                    num_workers=num_workers, **kwargs)
 
     old_x = data.x.clone().cpu()
-    new_datamodule_repr = has_package('pytorch_lightning>=2.5.0')
     flag = loader != 'full'
-    datamodule_repr = (
-        '{Train dataloader: ' + f'size={140 if flag else 1}' + '}\n'
-        '{Validation dataloader: ' + f'size={500 if flag else 1}' + '}\n'
-        '{Test dataloader: ' + f'size={1000 if flag else 1}' + '}\n'
-        '{Predict dataloader: ' + f'size={2708 if flag else 1}' +
-        '}' if new_datamodule_repr else f'LightningNodeData(data={data_repr}, '
-        f'loader={loader}, batch_size={batch_size}, '
-        f'num_workers={num_workers}, {kwargs_repr}'
-        f'pin_memory={flag}, '
-        f'persistent_workers={flag})')
+    if has_package('pytorch_lightning>=2.5.0'):
+        datamodule_repr = (
+            '{Train dataloader: ' + f'size={140 if flag else 1}' + '}\n'
+            '{Validation dataloader: ' + f'size={500 if flag else 1}' + '}\n'
+            '{Test dataloader: ' + f'size={1000 if flag else 1}' + '}\n'
+            '{Predict dataloader: ' + f'size={2708 if flag else 1}' + '}')
+    else:
+        datamodule_repr = (f'LightningNodeData(data={data_repr}, '
+                           f'loader={loader}, batch_size={batch_size}, '
+                           f'num_workers={num_workers}, {kwargs_repr}'
+                           f'pin_memory={flag}, '
+                           f'persistent_workers={flag})')
     assert str(datamodule) == datamodule_repr
 
     trainer.fit(model, datamodule)
diff --git a/test/metrics/test_link_pred_metric.py b/test/metrics/test_link_pred_metric.py
index f4bb5c1af..1dd30a352 100644
--- a/test/metrics/test_link_pred_metric.py
+++ b/test/metrics/test_link_pred_metric.py
@@ -54,6 +54,11 @@ def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
     expected = torch.tensor(values).mean()
     assert torch.allclose(out, expected)
 
+    # Test with `k > pred_mat.size(1)`:
+    metric.update(top_k_pred_mat[:, :k - 1], edge_label_index)
+    metric.compute()
+    metric.reset()
+
 
 def test_recall():
     pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
@@ -66,6 +71,11 @@ def test_recall():
 
     assert float(result) == pytest.approx(0.5 * (2 / 3 + 0.5))
 
+    # Test with `k > pred_mat.size(1)`:
+    metric.update(pred_mat[:, :1], edge_label_index)
+    metric.compute()
+    metric.reset()
+
 
 def test_f1():
     pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
@@ -77,6 +87,11 @@ def test_f1():
     result = metric.compute()
     assert float(result) == pytest.approx(0.6500)
 
+    # Test with `k > pred_mat.size(1)`:
+    metric.update(pred_mat[:, :1], edge_label_index)
+    metric.compute()
+    metric.reset()
+
 
 def test_map():
     pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
@@ -88,6 +103,11 @@ def test_map():
     result = metric.compute()
     assert float(result) == pytest.approx(0.6250)
 
+    # Test with `k > pred_mat.size(1)`:
+    metric.update(pred_mat[:, :1], edge_label_index)
+    metric.compute()
+    metric.reset()
+
 
 def test_ndcg():
     pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
@@ -100,6 +120,11 @@ def test_ndcg():
 
     assert float(result) == pytest.approx(0.6934264)
 
+    # Test with `k > pred_mat.size(1)`:
+    metric.update(pred_mat[:, :1], edge_label_index)
+    metric.compute()
+    metric.reset()
+
 
 def test_mrr():
     pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2], [0, 1]])
@@ -111,3 +136,8 @@ def test_mrr():
     result = metric.compute()
 
     assert float(result) == pytest.approx((1 + 0.5 + 0) / 3)
+
+    # Test with `k > pred_mat.size(1)`:
+    metric.update(pred_mat[:, :1], edge_label_index)
+    metric.compute()
+    metric.reset()