pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/metrics/link_pred.py

@@ -1,4 +1,5 @@
-from typing import Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -14,7 +15,143 @@ except Exception:
     BaseMetric = torch.nn.Module  # type: ignore
 
 
-class LinkPredMetric(BaseMetric):
+@dataclass(repr=False)
+class LinkPredMetricData:
+    pred_index_mat: Tensor
+    edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]]
+    edge_label_weight: Optional[Tensor] = None
+
+    def __post_init__(self) -> None:
+        # Filter all negative weights - they should not be used as ground-truth
+        if self.edge_label_weight is not None:
+            pos_mask = self.edge_label_weight > 0
+            self.edge_label_weight = self.edge_label_weight[pos_mask]
+            if isinstance(self.edge_label_index, Tensor):
+                self.edge_label_index = self.edge_label_index[:, pos_mask]
+            else:
+                self.edge_label_index = (
+                    self.edge_label_index[0][pos_mask],
+                    self.edge_label_index[1][pos_mask],
+                )
+
+    @property
+    def pred_rel_mat(self) -> Tensor:
+        r"""Returns a matrix indicating the relevance of the `k`-th prediction.
+        If :obj:`edge_label_weight` is not given, relevance will be denoted as
+        binary.
+        """
+        if hasattr(self, '_pred_rel_mat'):
+            return self._pred_rel_mat  # type: ignore
+
+        if self.edge_label_index[1].numel() == 0:
+            self._pred_rel_mat = torch.zeros_like(
+                self.pred_index_mat,
+                dtype=torch.bool if self.edge_label_weight is None else
+                torch.get_default_dtype(),
+            )
+            return self._pred_rel_mat
+
+        # Flatten both prediction and ground-truth indices, and determine
+        # overlaps afterwards via `torch.searchsorted`.
+        max_index = max(
+            self.pred_index_mat.max()
+            if self.pred_index_mat.numel() > 0 else 0,
+            self.edge_label_index[1].max()
+            if self.edge_label_index[1].numel() > 0 else 0,
+        ) + 1
+        arange = torch.arange(
+            start=0,
+            end=max_index * self.pred_index_mat.size(0),  # type: ignore
+            step=max_index,  # type: ignore
+            device=self.pred_index_mat.device,
+        ).view(-1, 1)
+        flat_pred_index = (self.pred_index_mat + arange).view(-1)
+        flat_label_index = max_index * self.edge_label_index[0]
+        flat_label_index = flat_label_index + self.edge_label_index[1]
+        flat_label_index, perm = flat_label_index.sort()
+        edge_label_weight = self.edge_label_weight
+        if edge_label_weight is not None:
+            assert edge_label_weight.size() == self.edge_label_index[0].size()
+            edge_label_weight = edge_label_weight[perm]
+
+        pos = torch.searchsorted(flat_label_index, flat_pred_index)
+        pos = pos.clamp(max=flat_label_index.size(0) - 1)  # Out-of-bounds.
+
+        pred_rel_mat = flat_label_index[pos] == flat_pred_index  # Find matches
+        if edge_label_weight is not None:
+            pred_rel_mat = edge_label_weight[pos].where(
+                pred_rel_mat,
+                pred_rel_mat.new_zeros(1),
+            )
+        pred_rel_mat = pred_rel_mat.view(self.pred_index_mat.size())
+
+        self._pred_rel_mat = pred_rel_mat
+        return pred_rel_mat
+
+    @property
+    def label_count(self) -> Tensor:
+        r"""The number of ground-truth labels for every example."""
+        if hasattr(self, '_label_count'):
+            return self._label_count  # type: ignore
+
+        label_count = scatter(
+            torch.ones_like(self.edge_label_index[0]),
+            self.edge_label_index[0],
+            dim=0,
+            dim_size=self.pred_index_mat.size(0),
+            reduce='sum',
+        )
+
+        self._label_count = label_count
+        return label_count
+
+    @property
+    def label_weight_sum(self) -> Tensor:
+        r"""The sum of edge label weights for every example."""
+        if self.edge_label_weight is None:
+            return self.label_count
+
+        if hasattr(self, '_label_weight_sum'):
+            return self._label_weight_sum  # type: ignore
+
+        label_weight_sum = scatter(
+            self.edge_label_weight,
+            self.edge_label_index[0],
+            dim=0,
+            dim_size=self.pred_index_mat.size(0),
+            reduce='sum',
+        )
+
+        self._label_weight_sum = label_weight_sum
+        return label_weight_sum
+
+    @property
+    def edge_label_weight_pos(self) -> Optional[Tensor]:
+        r"""Returns the position of edge label weights in descending order
+        within example-wise buckets.
+        """
+        if self.edge_label_weight is None:
+            return None
+
+        if hasattr(self, '_edge_label_weight_pos'):
+            return self._edge_label_weight_pos  # type: ignore
+
+        # Get the permutation via two sorts: One globally on the weights,
+        # followed by a (stable) sort on the example indices.
+        perm1 = self.edge_label_weight.argsort(descending=True)
+        perm2 = self.edge_label_index[0][perm1].argsort(stable=True)
+        perm = perm1[perm2]
+        # Invert the permutation to get the final position:
+        pos = torch.empty_like(perm)
+        pos[perm] = torch.arange(perm.size(0), device=perm.device)
+        # Normalize position to zero within all buckets:
+        pos = pos - cumsum(self.label_count)[self.edge_label_index[0]]
+
+        self._edge_label_weight_pos = pos
+        return pos
+
+
+class _LinkPredMetric(BaseMetric):
     r"""An abstract class for computing link prediction retrieval metrics.
 
     Args:
@@ -33,20 +170,11 @@ class LinkPredMetric(BaseMetric):
 
         self.k = k
 
-        self.accum: Tensor
-        self.total: Tensor
-
-        if WITH_TORCHMETRICS:
-            self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
-            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
-        else:
-            self.register_buffer('accum', torch.tensor(0.))
-            self.register_buffer('total', torch.tensor(0))
-
     def update(
         self,
         pred_index_mat: Tensor,
         edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
     ) -> None:
         r"""Updates the state variables based on the current mini-batch
         prediction.
@@ -62,99 +190,293 @@ class LinkPredMetric(BaseMetric):
             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]`.
+            edge_label_weight (torch.Tensor, optional): The weight of the
+                ground-truth indices for every example in the mini-batch of
+                shape :obj:`[num_ground_truth_indices]`. If given, needs to be
+                a vector of positive values. Required for weighted metrics,
+                ignored otherwise. (default: :obj:`None`)
         """
-        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`.
-        max_index = max(  # type: ignore
-            pred_index_mat.max() if pred_index_mat.numel() > 0 else 0,
-            edge_label_index[1].max()
-            if edge_label_index[1].numel() > 0 else 0,
-        ) + 1
-        arange = torch.arange(
-            start=0,
-            end=max_index * pred_index_mat.size(0),
-            step=max_index,
-            device=pred_index_mat.device,
-        ).view(-1, 1)
-        flat_pred_index = (pred_index_mat + arange).view(-1)
-        flat_y_index = max_index * edge_label_index[0] + edge_label_index[1]
+        raise NotImplementedError
 
-        pred_isin_mat = torch.isin(flat_pred_index, flat_y_index)
-        pred_isin_mat = pred_isin_mat.view(pred_index_mat.size())
+    def compute(self) -> Tensor:
+        r"""Computes the final metric value."""
+        raise NotImplementedError
 
-        # Compute the number of targets per example:
-        y_count = scatter(
-            torch.ones_like(edge_label_index[0]),
-            edge_label_index[0],
-            dim=0,
-            dim_size=pred_index_mat.size(0),
-            reduce='sum',
+    def reset(self) -> None:
+        r"""Resets metric state variables to their default value."""
+        if WITH_TORCHMETRICS:
+            super().reset()
+        else:
+            self._reset()
+
+    def _reset(self) -> None:
+        raise NotImplementedError
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(k={self.k})'
+
+
+class LinkPredMetric(_LinkPredMetric):
+    r"""An abstract class for computing link prediction retrieval metrics.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    weighted: bool
+
+    def __init__(self, k: int) -> None:
+        super().__init__(k)
+
+        self.accum: Tensor
+        self.total: Tensor
+
+        if WITH_TORCHMETRICS:
+            self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
+        else:
+            self.register_buffer('accum', torch.tensor(0.), persistent=False)
+            self.register_buffer('total', torch.tensor(0), persistent=False)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        if self.weighted and edge_label_weight is None:
+            raise ValueError(f"'edge_label_weight' is a required argument for "
+                             f"weighted '{self.__class__.__name__}' metrics")
+        if not self.weighted:
+            edge_label_weight = None
+
+        data = LinkPredMetricData(
+            pred_index_mat=pred_index_mat,
+            edge_label_index=edge_label_index,
+            edge_label_weight=edge_label_weight,
         )
+        self._update(data)
 
-        metric = self._compute(pred_isin_mat, y_count)
+    def _update(self, data: LinkPredMetricData) -> None:
+        metric = self._compute(data)
 
         self.accum += metric.sum()
-        self.total += (y_count > 0).sum()
+        self.total += (data.label_count > 0).sum()
 
     def compute(self) -> Tensor:
-        r"""Computes the final metric value."""
         if self.total == 0:
             return torch.zeros_like(self.accum)
         return self.accum / self.total
 
-    def reset(self) -> None:
-        r"""Reset metric state variables to their default value."""
-        if WITH_TORCHMETRICS:
-            super().reset()
-        else:
-            self.accum.zero_()
-            self.total.zero_()
-
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        r"""Compute the specific metric.
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        r"""Computes the specific metric.
         To be implemented separately for each metric class.
 
         Args:
-            pred_isin_mat (torch.Tensor): A boolean matrix whose :obj:`(i,k)`
-                element indicates if the :obj:`k`-th prediction for the
-                :obj:`i`-th example is correct or not.
-            y_count (torch.Tensor): A vector indicating the number of
-                ground-truth labels for each example.
+            data (LinkPredMetricData): The mini-batch data for computing a link
+                prediction metric per example.
         """
         raise NotImplementedError
 
+    def _reset(self) -> None:
+        self.accum.zero_()
+        self.total.zero_()
+
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}(k={self.k})'
+        weighted_repr = ', weighted=True' if self.weighted else ''
+        return f'{self.__class__.__name__}(k={self.k}{weighted_repr})'
+
+
+class LinkPredMetricCollection(torch.nn.ModuleDict):
+    r"""A collection of metrics to reduce and speed-up computation of link
+    prediction metrics.
+
+    .. code-block:: python
+
+        from torch_geometric.metrics import (
+            LinkPredMAP,
+            LinkPredMetricCollection,
+            LinkPredPrecision,
+            LinkPredRecall,
+        )
+
+        metrics = LinkPredMetricCollection([
+            LinkPredMAP(k=10),
+            LinkPredPrecision(k=100),
+            LinkPredRecall(k=50),
+        ])
+
+        metrics.update(pred_index_mat, edge_label_index)
+        out = metrics.compute()
+        metrics.reset()
+
+        print(out)
+        >>> {'LinkPredMAP@10': tensor(0.375),
+        ...  'LinkPredPrecision@100': tensor(0.127),
+        ...  'LinkPredRecall@50': tensor(0.483)}
+
+    Args:
+        metrics: The link prediction metrics.
+    """
+    def __init__(
+        self,
+        metrics: Union[
+            List[LinkPredMetric],
+            Dict[str, LinkPredMetric],
+        ],
+    ) -> None:
+        super().__init__()
+
+        if isinstance(metrics, (list, tuple)):
+            metrics = {
+                (f'{"Weighted" if getattr(metric, "weighted", False) else ""}'
+                 f'{metric.__class__.__name__}@{metric.k}'):
+                metric
+                for metric in metrics
+            }
+        assert len(metrics) > 0
+        assert isinstance(metrics, dict)
+
+        for name, metric in metrics.items():
+            assert isinstance(metric, _LinkPredMetric)
+            self[name] = metric
+
+    @property
+    def max_k(self) -> int:
+        r"""The maximum number of top-:math:`k` predictions to evaluate
+        against.
+        """
+        return max([
+            metric.k  # type: ignore[return-value]
+            for metric in self.values()
+        ])  # type: ignore[type-var]
+
+    @property
+    def weighted(self) -> bool:
+        r"""Returns :obj:`True` in case the collection holds at least one
+        weighted link prediction metric.
+        """
+        return any(
+            [getattr(metric, 'weighted', False) for metric in self.values()])
+
+    def update(  # type: ignore
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> 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]`.
+            edge_label_weight (torch.Tensor, optional): The weight of the
+                ground-truth indices for every example in the mini-batch of
+                shape :obj:`[num_ground_truth_indices]`. If given, needs to be
+                a vector of positive values. Required for weighted metrics,
+                ignored otherwise. (default: :obj:`None`)
+        """
+        if self.weighted and edge_label_weight is None:
+            raise ValueError(f"'edge_label_weight' is a required argument for "
+                             f"weighted '{self.__class__.__name__}' metrics")
+
+        data = LinkPredMetricData(  # Share metric data across metrics.
+            pred_index_mat=pred_index_mat,
+            edge_label_index=edge_label_index,
+            edge_label_weight=edge_label_weight,
+        )
+
+        for metric in self.values():
+            if isinstance(metric, LinkPredMetric) and metric.weighted:
+                metric._update(data)
+                if WITH_TORCHMETRICS:
+                    metric._update_count += 1
+
+        data.edge_label_weight = None
+        if hasattr(data, '_pred_rel_mat'):
+            data._pred_rel_mat = data._pred_rel_mat != 0.0
+        if hasattr(data, '_label_weight_sum'):
+            del data._label_weight_sum
+        if hasattr(data, '_edge_label_weight_pos'):
+            del data._edge_label_weight_pos
+
+        for metric in self.values():
+            if isinstance(metric, LinkPredMetric) and not metric.weighted:
+                metric._update(data)
+                if WITH_TORCHMETRICS:
+                    metric._update_count += 1
+
+        for metric in self.values():
+            if not isinstance(metric, LinkPredMetric):
+                metric.update(  # type: ignore[operator]
+                    pred_index_mat,
+                    edge_label_index,
+                    edge_label_weight,
+                )
+
+    def compute(self) -> Dict[str, Tensor]:
+        r"""Computes the final metric values."""
+        return {
+            name: metric.compute()  # type: ignore[operator]
+            for name, metric in self.items()
+        }
+
+    def reset(self) -> None:
+        r"""Reset metric state variables to their default value."""
+        for metric in self.values():
+            metric.reset()  # type: ignore[operator]
+
+    def __repr__(self) -> str:
+        names = [f'  {name}: {metric},\n' for name, metric in self.items()]
+        return f'{self.__class__.__name__}([\n{"".join(names)}])'
 
 
 class LinkPredPrecision(LinkPredMetric):
-    r"""A link prediction metric to compute Precision @ :math:`k`.
+    r"""A link prediction metric to compute Precision @ :math:`k`, *i.e.* the
+    proportion of recommendations within the top-:math:`k` that are actually
+    relevant.
+
+    A higher precision indicates the model's ability to surface relevant items
+    early in the ranking.
 
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        return pred_isin_mat.sum(dim=-1) / self.k
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        return pred_rel_mat.sum(dim=-1) / self.k
 
 
 class LinkPredRecall(LinkPredMetric):
-    r"""A link prediction metric to compute Recall @ :math:`k`.
+    r"""A link prediction metric to compute Recall @ :math:`k`, *i.e.* the
+    proportion of relevant items that appear within the top-:math:`k`.
+
+    A higher recall indicates the model's ability to retrieve a larger
+    proportion of relevant items.
 
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        return pred_isin_mat.sum(dim=-1) / y_count.clamp(min=1e-7)
+    def __init__(self, k: int, weighted: bool = False):
+        super().__init__(k=k)
+        self.weighted = weighted
+
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        return pred_rel_mat.sum(dim=-1) / data.label_weight_sum.clamp(min=1e-7)
 
 
 class LinkPredF1(LinkPredMetric):
@@ -164,55 +486,403 @@ class LinkPredF1(LinkPredMetric):
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        isin_count = pred_isin_mat.sum(dim=-1)
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        isin_count = pred_rel_mat.sum(dim=-1)
         precision = isin_count / self.k
-        recall = isin_count = isin_count / y_count.clamp(min=1e-7)
+        recall = isin_count / data.label_count.clamp(min=1e-7)
         return 2 * precision * recall / (precision + recall).clamp(min=1e-7)
 
 
 class LinkPredMAP(LinkPredMetric):
     r"""A link prediction metric to compute MAP @ :math:`k` (Mean Average
-    Precision).
+    Precision), considering the order of relevant items within the
+    top-:math:`k`.
+
+    MAP @ :math:`k` can provide a more comprehensive view of ranking quality
+    than precision alone.
 
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    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))
-        return ((cum_precision * pred_isin_mat).sum(dim=-1) /
-                y_count.clamp(min=1e-7, max=self.k))
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        device = pred_rel_mat.device
+        arange = torch.arange(1, pred_rel_mat.size(1) + 1, device=device)
+        cum_precision = pred_rel_mat.cumsum(dim=1) / arange
+        return ((cum_precision * pred_rel_mat).sum(dim=-1) /
+                data.label_count.clamp(min=1e-7, max=self.k))
 
 
 class LinkPredNDCG(LinkPredMetric):
     r"""A link prediction metric to compute the NDCG @ :math:`k` (Normalized
     Discounted Cumulative Gain).
 
+    In particular, can account for the position of relevant items by
+    considering relevance scores, giving higher weight to more relevant items
+    appearing at the top.
+
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
+        weighted (bool, optional): If set to :obj:`True`, assumes sorted lists
+            of ground-truth items according to a relevance score as given by
+            :obj:`edge_label_weight`. (default: :obj:`False`)
     """
     higher_is_better: bool = True
 
-    def __init__(self, k: int):
+    def __init__(self, k: int, weighted: bool = False):
         super().__init__(k=k)
+        self.weighted = weighted
 
         dtype = torch.get_default_dtype()
-        multiplier = 1.0 / torch.arange(2, k + 2, dtype=dtype).log2()
+        discount = torch.arange(2, k + 2, dtype=dtype).log2()
+
+        self.discount: Tensor
+        self.register_buffer('discount', discount, persistent=False)
 
-        self.multiplier: Tensor
-        self.register_buffer('multiplier', multiplier)
+        if not weighted:
+            self.register_buffer('idcg', cumsum(1.0 / discount),
+                                 persistent=False)
+        else:
+            self.idcg = None
 
-        self.idcg: Tensor
-        self.register_buffer('idcg', cumsum(multiplier))
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        discount = self.discount[:pred_rel_mat.size(1)].view(1, -1)
+        dcg = (pred_rel_mat / discount).sum(dim=-1)
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        dcg = (pred_isin_mat * self.multiplier.view(1, -1)).sum(dim=-1)
-        idcg = self.idcg[y_count.clamp(max=self.k)]
+        if not self.weighted:
+            assert self.idcg is not None
+            idcg = self.idcg[data.label_count.clamp(max=self.k)]
+        else:
+            assert data.edge_label_weight is not None
+            pos = data.edge_label_weight_pos
+            assert pos is not None
+
+            discount = torch.cat([
+                self.discount,
+                self.discount.new_full((1, ), fill_value=float('inf')),
+            ])
+            discount = discount[pos.clamp(max=self.k)]
+
+            idcg = scatter(  # Apply discount and aggregate:
+                data.edge_label_weight / discount,
+                data.edge_label_index[0],
+                dim_size=data.pred_index_mat.size(0),
+                reduce='sum',
+            )
 
         out = dcg / idcg
         out[out.isnan() | out.isinf()] = 0.0
         return out
+
+
+class LinkPredMRR(LinkPredMetric):
+    r"""A link prediction metric to compute the MRR @ :math:`k` (Mean
+    Reciprocal Rank), *i.e.* the mean reciprocal rank of the first correct
+    prediction (or zero otherwise).
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    higher_is_better: bool = True
+    weighted: bool = False
+
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        device = pred_rel_mat.device
+        arange = torch.arange(1, pred_rel_mat.size(1) + 1, device=device)
+        return (pred_rel_mat / arange).max(dim=-1)[0]
+
+
+class LinkPredHitRatio(LinkPredMetric):
+    r"""A link prediction metric to compute the hit ratio @ :math:`k`, *i.e.*
+    the percentage of users for whom at least one relevant item is present
+    within the top-:math:`k` recommendations.
+
+    A high ratio signifies the model's effectiveness in satisfying a broad
+    range of user preferences.
+    """
+    higher_is_better: bool = True
+    weighted: bool = False
+
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        return pred_rel_mat.max(dim=-1)[0].to(torch.get_default_dtype())
+
+
+class LinkPredCoverage(_LinkPredMetric):
+    r"""A link prediction metric to compute the Coverage @ :math:`k` of
+    predictions, *i.e.* the percentage of unique items recommended across all
+    users within the top-:math:`k`.
+
+    Higher coverage indicates a wider exploration of the item catalog.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+        num_dst_nodes (int): The total number of destination nodes.
+    """
+    higher_is_better: bool = True
+
+    def __init__(self, k: int, num_dst_nodes: int) -> None:
+        super().__init__(k)
+        self.num_dst_nodes = num_dst_nodes
+
+        self.mask: Tensor
+        mask = torch.zeros(num_dst_nodes, dtype=torch.bool)
+        if WITH_TORCHMETRICS:
+            self.add_state('mask', mask, dist_reduce_fx='max')
+        else:
+            self.register_buffer('mask', mask, persistent=False)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        self.mask[pred_index_mat[:, :self.k].flatten()] = True
+
+    def compute(self) -> Tensor:
+        return self.mask.to(torch.get_default_dtype()).mean()
+
+    def _reset(self) -> None:
+        self.mask.zero_()
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(k={self.k}, '
+                f'num_dst_nodes={self.num_dst_nodes})')
+
+
+class LinkPredDiversity(_LinkPredMetric):
+    r"""A link prediction metric to compute the Diversity @ :math:`k` of
+    predictions according to item categories.
+
+    Diversity is computed as
+
+    .. math::
+        div_{u@k} = 1 - \left( \frac{1}{k \cdot (k-1)} \right) \sum_{i \neq j}
+        sim(i, j)
+
+    where
+
+    .. math::
+        sim(i,j) = \begin{cases}
+            1 & \quad \text{if } i,j \text{ share category,}\\
+            0 & \quad \text{otherwise.}
+        \end{cases}
+
+    which measures the pair-wise inequality of recommendations according to
+    item categories.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+        category (torch.Tensor): A vector that assigns each destination node to
+            a specific category.
+    """
+    higher_is_better: bool = True
+
+    def __init__(self, k: int, category: Tensor) -> None:
+        super().__init__(k)
+
+        self.accum: Tensor
+        self.total: Tensor
+
+        if WITH_TORCHMETRICS:
+            self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
+        else:
+            self.register_buffer('accum', torch.tensor(0.), persistent=False)
+            self.register_buffer('total', torch.tensor(0), persistent=False)
+
+        self.category: Tensor
+        self.register_buffer('category', category, persistent=False)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        category = self.category[pred_index_mat[:, :self.k]]
+
+        sim = (category.unsqueeze(-2) == category.unsqueeze(-1)).sum(dim=-1)
+        div = 1 - 1 / (self.k * (self.k - 1)) * (sim - 1).sum(dim=-1)
+
+        self.accum += div.sum()
+        self.total += pred_index_mat.size(0)
+
+    def compute(self) -> Tensor:
+        if self.total == 0:
+            return torch.zeros_like(self.accum)
+        return self.accum / self.total
+
+    def _reset(self) -> None:
+        self.accum.zero_()
+        self.total.zero_()
+
+
+class LinkPredPersonalization(_LinkPredMetric):
+    r"""A link prediction metric to compute the Personalization @ :math:`k`,
+    *i.e.* the dissimilarity of recommendations across different users.
+
+    Higher personalization suggests that the model tailors recommendations to
+    individual user preferences rather than providing generic results.
+
+    Dissimilarity is defined by the average inverse cosine similarity between
+    users' lists of recommendations.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+        max_src_nodes (int, optional): The maximum source nodes to consider to
+            compute pair-wise dissimilarity. If specified,
+            Personalization @ :math:`k` is approximated to avoid computation
+            blowup due to quadratic complexity. (default: :obj:`2**12`)
+        batch_size (int, optional): The batch size to determine how many pairs
+            of user recommendations should be processed at once.
+            (default: :obj:`2**16`)
+    """
+    higher_is_better: bool = True
+
+    def __init__(
+        self,
+        k: int,
+        max_src_nodes: Optional[int] = 2**12,
+        batch_size: int = 2**16,
+    ) -> None:
+        super().__init__(k)
+        self.max_src_nodes = max_src_nodes
+        self.batch_size = batch_size
+
+        self.preds: List[Tensor]
+        self.total: Tensor
+
+        if WITH_TORCHMETRICS:
+            self.add_state('preds', default=[], dist_reduce_fx='cat')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
+        else:
+            self.preds = []
+            self.register_buffer('total', torch.tensor(0), persistent=False)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+
+        # NOTE Move to CPU to avoid memory blowup.
+        pred_index_mat = pred_index_mat[:, :self.k].cpu()
+
+        if self.max_src_nodes is None:
+            self.preds.append(pred_index_mat)
+            self.total += pred_index_mat.size(0)
+        elif self.total < self.max_src_nodes:
+            remaining = int(self.max_src_nodes - self.total)
+            pred_index_mat = pred_index_mat[:remaining]
+            self.preds.append(pred_index_mat)
+            self.total += pred_index_mat.size(0)
+
+    def compute(self) -> Tensor:
+        device = self.total.device
+        score = torch.tensor(0.0, device=device)
+        total = torch.tensor(0, device=device)
+
+        if len(self.preds) == 0:
+            return score
+
+        pred = torch.cat(self.preds, dim=0)
+
+        if pred.size(0) == 0:
+            return score
+
+        # Calculate all pairs of nodes (e.g., triu_indices with offset=1).
+        # NOTE We do this in chunks to avoid memory blow-up, which leads to a
+        # more efficient but trickier implementation.
+        num_pairs = (pred.size(0) * (pred.size(0) - 1)) // 2
+        offset = torch.arange(pred.size(0) - 1, 0, -1, device=device)
+        rowptr = cumsum(offset)
+        for start in range(0, num_pairs, self.batch_size):
+            end = min(start + self.batch_size, num_pairs)
+            idx = torch.arange(start, end, device=device)
+
+            # Find the corresponding row:
+            row = torch.searchsorted(rowptr, idx, right=True) - 1
+            # Find the corresponding column:
+            col = idx - rowptr[row] + (pred.size(0) - offset[row])
+
+            left = pred[row.cpu()].to(device)
+            right = pred[col.cpu()].to(device)
+
+            # Use offset to work around applying `isin` along a specific dim:
+            i = max(int(left.max()), int(right.max())) + 1
+            idx = torch.arange(0, i * row.size(0), i, device=device)
+            idx = idx.view(-1, 1)
+            isin = torch.isin(left + idx, right + idx)
+
+            # Compute personalization via average inverse cosine similarity:
+            cos = isin.sum(dim=-1) / pred.size(1)
+            score += (1 - cos).sum()
+            total += cos.numel()
+
+        return score / total
+
+    def _reset(self) -> None:
+        self.preds = []
+        self.total.zero_()
+
+
+class LinkPredAveragePopularity(_LinkPredMetric):
+    r"""A link prediction metric to compute the Average Recommendation
+    Popularity (ARP) @ :math:`k`, which provides insights into the model's
+    tendency to recommend popular items by averaging the popularity scores of
+    items within the top-:math:`k` recommendations.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+        popularity (torch.Tensor): The popularity of every item in the training
+            set, *e.g.*, the number of times an item has been rated.
+    """
+    higher_is_better: bool = False
+
+    def __init__(self, k: int, popularity: Tensor) -> None:
+        super().__init__(k)
+
+        self.accum: Tensor
+        self.total: Tensor
+
+        if WITH_TORCHMETRICS:
+            self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
+        else:
+            self.register_buffer('accum', torch.tensor(0.), persistent=False)
+            self.register_buffer('total', torch.tensor(0), persistent=False)
+
+        self.popularity: Tensor
+        self.register_buffer('popularity', popularity, persistent=False)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        pred_index_mat = pred_index_mat[:, :self.k]
+        popularity = self.popularity[pred_index_mat]
+        popularity = popularity.to(self.accum.dtype).mean(dim=-1)
+        self.accum += popularity.sum()
+        self.total += popularity.numel()
+
+    def compute(self) -> Tensor:
+        if self.total == 0:
+            return torch.zeros_like(self.accum)
+        return self.accum / self.total
+
+    def _reset(self) -> None:
+        self.accum.zero_()
+        self.total.zero_()