pyg-nightly 2.6.0.dev20240319__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/metrics/link_pred.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -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()
@@ -78,8 +59,8 @@ class LinkPredMetric(BaseMetric):
         ) + 1
         arange = torch.arange(
             start=0,
-            end=max_index * pred_index_mat.size(0),
-            step=max_index,
+            end=max_index * pred_index_mat.size(0),  # type: ignore
+            step=max_index,  # type: ignore
             device=pred_index_mat.device,
         ).view(-1, 1)
         flat_pred_index = (pred_index_mat + arange).view(-1)
@@ -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:
@@ -133,6 +144,103 @@ class LinkPredMetric(BaseMetric):
         return f'{self.__class__.__name__}(k={self.k})'
 
 
+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'{metric.__class__.__name__}@{metric.k}': metric
+                for metric in metrics
+            }
+        assert len(metrics) > 0
+        assert isinstance(metrics, dict)
+
+        for name, metric in metrics.items():
+            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 for metric in self.values()])
+
+    def update(  # type: ignore
+        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 = LinkPredMetric._prepare(
+            pred_index_mat, edge_label_index)
+        for metric in self.values():
+            metric._update_from_prepared(pred_isin_mat, y_count)
+
+    def compute(self) -> Dict[str, Tensor]:
+        r"""Computes the final metric values."""
+        return {name: metric.compute() 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()
+
+    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`.
 
@@ -182,8 +290,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,9 +319,25 @@ 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
         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).
+
+    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:
+        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]

torch_geometric/nn/aggr/__init__.py

@@ -24,6 +24,8 @@ from .mlp import MLPAggregation
 from .deep_sets import DeepSetsAggregation
 from .set_transformer import SetTransformerAggregation
 from .lcm import LCMAggregation
+from .variance_preserving import VariancePreservingAggregation
+from .patch_transformer import PatchTransformerAggregation
 
 __all__ = classes = [
     'Aggregation',
@@ -51,4 +53,6 @@ __all__ = classes = [
     'DeepSetsAggregation',
     'SetTransformerAggregation',
     'LCMAggregation',
+    'VariancePreservingAggregation',
+    'PatchTransformerAggregation',
 ]

torch_geometric/nn/aggr/attention.py

@@ -65,8 +65,6 @@ class AttentionalAggregation(Aggregation):
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
 
-        self.assert_two_dimensional_input(x, dim)
-
         if self.gate_mlp is not None:
             gate = self.gate_mlp(x, batch=index, batch_size=dim_size)
         else:

torch_geometric/nn/aggr/base.py

@@ -25,7 +25,7 @@ class Aggregation(torch.nn.Module):
     Notably, :obj:`index` does not have to be sorted (for most aggregation
     operators):
 
-    .. code-block::
+    .. code-block:: python
 
        # Feature matrix holding 10 elements with 64 features each:
        x = torch.randn(10, 64)
@@ -39,7 +39,7 @@ class Aggregation(torch.nn.Module):
     called :obj:`ptr`. Here, elements within the same set need to be grouped
     together in the input, and :obj:`ptr` defines their boundaries:
 
-    .. code-block::
+    .. code-block:: python
 
        # Feature matrix holding 10 elements with 64 features each:
        x = torch.randn(10, 64)
@@ -94,11 +94,9 @@ class Aggregation(torch.nn.Module):
             max_num_elements: (int, optional): The maximum number of elements
                 within a single aggregation group. (default: :obj:`None`)
         """
-        pass
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        pass
 
     @disable_dynamic_shapes(required_args=['dim_size'])
     def __call__(

torch_geometric/nn/aggr/patch_transformer.py

@@ -0,0 +1,143 @@
+import math
+from typing import List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.experimental import disable_dynamic_shapes
+from torch_geometric.nn.aggr import Aggregation
+from torch_geometric.nn.aggr.utils import MultiheadAttentionBlock
+from torch_geometric.nn.encoding import PositionalEncoding
+from torch_geometric.utils import scatter
+
+
+class PatchTransformerAggregation(Aggregation):
+    r"""Performs patch transformer aggregation in which the elements to
+    aggregate are processed by multi-head attention blocks across patches, as
+    described in the `"Simplifying Temporal Heterogeneous Network for
+    Continuous-Time Link Prediction"
+    <https://dl.acm.org/doi/pdf/10.1145/3583780.3615059>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        patch_size (int): Number of elements in a patch.
+        hidden_channels (int): Intermediate size of each sample.
+        num_transformer_blocks (int, optional): Number of transformer blocks
+            (default: :obj:`1`).
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        dropout (float, optional): Dropout probability of attention weights.
+            (default: :obj:`0.0`)
+        aggr (str or list[str], optional): The aggregation module, *e.g.*,
+            :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
+            :obj:`"var"`, :obj:`"std"`. (default: :obj:`"mean"`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        patch_size: int,
+        hidden_channels: int,
+        num_transformer_blocks: int = 1,
+        heads: int = 1,
+        dropout: float = 0.0,
+        aggr: Union[str, List[str]] = 'mean',
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.patch_size = patch_size
+        self.aggrs = [aggr] if isinstance(aggr, str) else aggr
+
+        assert len(self.aggrs) > 0
+        for aggr in self.aggrs:
+            assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
+
+        self.lin = torch.nn.Linear(in_channels, hidden_channels)
+        self.pad_projector = torch.nn.Linear(
+            patch_size * hidden_channels,
+            hidden_channels,
+        )
+        self.pe = PositionalEncoding(hidden_channels)
+
+        self.blocks = torch.nn.ModuleList([
+            MultiheadAttentionBlock(
+                channels=hidden_channels,
+                heads=heads,
+                layer_norm=True,
+                dropout=dropout,
+            ) for _ in range(num_transformer_blocks)
+        ])
+
+        self.fc = torch.nn.Linear(
+            hidden_channels * len(self.aggrs),
+            out_channels,
+        )
+
+    def reset_parameters(self) -> None:
+        self.lin.reset_parameters()
+        self.pad_projector.reset_parameters()
+        self.pe.reset_parameters()
+        for block in self.blocks:
+            block.reset_parameters()
+        self.fc.reset_parameters()
+
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
+    def forward(
+        self,
+        x: Tensor,
+        index: Tensor,
+        ptr: Optional[Tensor] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+    ) -> Tensor:
+
+        if max_num_elements is None:
+            if ptr is not None:
+                count = ptr.diff()
+            else:
+                count = scatter(torch.ones_like(index), index, dim=0,
+                                dim_size=dim_size, reduce='sum')
+            max_num_elements = int(count.max()) + 1
+
+        # Set `max_num_elements` to a multiple of `patch_size`:
+        max_num_elements = (math.floor(max_num_elements / self.patch_size) *
+                            self.patch_size)
+
+        x = self.lin(x)
+
+        # TODO If groups are heavily unbalanced, this will create a lot of
+        # "empty" patches. Try to figure out a way to fix this.
+        # [batch_size, num_patches * patch_size, hidden_channels]
+        x, _ = self.to_dense_batch(x, index, ptr, dim_size, dim,
+                                   max_num_elements=max_num_elements)
+
+        # [batch_size, num_patches, patch_size * hidden_channels]
+        x = x.view(x.size(0), max_num_elements // self.patch_size,
+                   self.patch_size * x.size(-1))
+
+        # [batch_size, num_patches, hidden_channels]
+        x = self.pad_projector(x)
+
+        x = x + self.pe(torch.arange(x.size(1), device=x.device))
+
+        # [batch_size, num_patches, hidden_channels]
+        for block in self.blocks:
+            x = block(x, x)
+
+        # [batch_size, hidden_channels]
+        outs: List[Tensor] = []
+        for aggr in self.aggrs:
+            out = getattr(torch, aggr)(x, dim=1)
+            outs.append(out[0] if isinstance(out, tuple) else out)
+        out = torch.cat(outs, dim=1) if len(outs) > 1 else outs[0]
+
+        # [batch_size, out_channels]
+        return self.fc(out)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, patch_size={self.patch_size})')

torch_geometric/nn/aggr/set_transformer.py

@@ -38,7 +38,7 @@ class SetTransformerAggregation(Aggregation):
             (default: :obj:`1`)
         concat (bool, optional): If set to :obj:`False`, the seed embeddings
             are averaged instead of concatenated. (default: :obj:`True`)
-        norm (str, optional): If set to :obj:`True`, will apply layer
+        layer_norm (str, optional): If set to :obj:`True`, will apply layer
             normalization. (default: :obj:`False`)
         dropout (float, optional): Dropout probability of attention weights.
             (default: :obj:`0`)

torch_geometric/nn/aggr/variance_preserving.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+from torch import Tensor
+
+from torch_geometric.nn.aggr import Aggregation
+from torch_geometric.utils import degree
+from torch_geometric.utils._scatter import broadcast
+
+
+class VariancePreservingAggregation(Aggregation):
+    r"""Performs the Variance Preserving Aggregation (VPA) from the `"GNN-VPA:
+    A Variance-Preserving Aggregation Strategy for Graph Neural Networks"
+    <https://arxiv.org/abs/2403.04747>`_ paper.
+
+    .. math::
+        \mathrm{vpa}(\mathcal{X}) = \frac{1}{\sqrt{|\mathcal{X}|}}
+        \sum_{\mathbf{x}_i \in \mathcal{X}} \mathbf{x}_i
+    """
+    def forward(self, x: Tensor, index: Optional[Tensor] = None,
+                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
+                dim: int = -2) -> Tensor:
+
+        out = self.reduce(x, index, ptr, dim_size, dim, reduce='sum')
+
+        if ptr is not None:
+            count = ptr.diff().to(out.dtype)
+        else:
+            count = degree(index, dim_size, dtype=out.dtype)
+
+        count = count.sqrt().clamp(min=1.0)
+        count = broadcast(count, ref=out, dim=dim)
+
+        return out / count

torch_geometric/nn/attention/__init__.py

@@ -1,3 +1,7 @@
 from .performer import PerformerAttention
+from .qformer import QFormer
 
-__all__ = ['PerformerAttention']
+__all__ = [
+    'PerformerAttention',
+    'QFormer',
+]

torch_geometric/nn/attention/qformer.py

@@ -0,0 +1,71 @@
+from typing import Callable
+
+import torch
+
+
+class QFormer(torch.nn.Module):
+    r"""The Querying Transformer (Q-Former) from
+    `"BLIP-2: Bootstrapping Language-Image Pre-training
+    with Frozen Image Encoders and Large Language Models"
+    <https://arxiv.org/pdf/2301.12597>`_ paper.
+
+    Args:
+        input_dim (int): The number of features in the input.
+        hidden_dim (int): The dimension of the fnn in the encoder layer.
+        output_dim (int): The final output dimension.
+        num_heads (int): The number of multi-attention-heads.
+        num_layers (int): The number of sub-encoder-layers in the encoder.
+        dropout (int): The dropout value in each encoder layer.
+
+
+    .. note::
+        This is a simplified version of the original Q-Former implementation.
+    """
+    def __init__(
+            self,
+            input_dim: int,
+            hidden_dim: int,
+            output_dim: int,
+            num_heads: int,
+            num_layers: int,
+            dropout: float = 0.0,
+            activation: Callable = torch.nn.ReLU(),
+    ) -> None:
+
+        super().__init__()
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+
+        self.layer_norm = torch.nn.LayerNorm(input_dim)
+        self.encoder_layer = torch.nn.TransformerEncoderLayer(
+            d_model=input_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_dim,
+            dropout=dropout,
+            activation=activation,
+            batch_first=True,
+        )
+        self.encoder = torch.nn.TransformerEncoder(
+            self.encoder_layer,
+            num_layers=num_layers,
+        )
+        self.project = torch.nn.Linear(input_dim, output_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): Input sequence to the encoder layer.
+                :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
+                batch-size :math:`B`, sequence length :math:`N`,
+                and feature dimension :math:`F`.
+        """
+        x = self.layer_norm(x)
+        x = self.encoder(x)
+        out = self.project(x)
+        return out
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}('
+                f'num_heads={self.num_heads}, '
+                f'num_layers={self.num_layers})')

torch_geometric/nn/conv/collect.jinja

@@ -4,8 +4,8 @@ import torch
 from torch import Tensor
 
 from torch_geometric import EdgeIndex
+from torch_geometric.index import ptr2index
 from torch_geometric.utils import is_torch_sparse_tensor
-from torch_geometric.utils.sparse import ptr2index
 from torch_geometric.typing import SparseTensor
 
 
@@ -98,13 +98,16 @@ def {{collect_name}}(
 
 {%- if 'edge_weight' in collect_param_dict and
     collect_param_dict['edge_weight'].type_repr.endswith('Tensor') %}
-    assert edge_weight is not None
+    if torch.jit.is_scripting():
+        assert edge_weight is not None
 {%- elif 'edge_attr' in collect_param_dict and
     collect_param_dict['edge_attr'].type_repr.endswith('Tensor') %}
-    assert edge_attr is not None
+    if torch.jit.is_scripting():
+        assert edge_attr is not None
 {%- elif 'edge_type' in collect_param_dict and
     collect_param_dict['edge_type'].type_repr.endswith('Tensor') %}
-    assert edge_type is not None
+    if torch.jit.is_scripting():
+        assert edge_type is not None
 {%- endif %}
 
     # Collect user-defined arguments:

torch_geometric/nn/conv/cugraph/base.py

@@ -7,12 +7,7 @@ from torch_geometric import EdgeIndex
 
 try:  # pragma: no cover
     LEGACY_MODE = False
-    from pylibcugraphops.pytorch import (
-        SampledCSC,
-        SampledHeteroCSC,
-        StaticCSC,
-        StaticHeteroCSC,
-    )
+    from pylibcugraphops.pytorch import CSC, HeteroCSC
     HAS_PYLIBCUGRAPHOPS = True
 except ImportError:
     HAS_PYLIBCUGRAPHOPS = False
@@ -41,7 +36,6 @@ class CuGraphModule(torch.nn.Module):  # pragma: no cover
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        pass
 
     def get_cugraph(
         self,
@@ -79,12 +73,13 @@ class CuGraphModule(torch.nn.Module):  # pragma: no cover
                 return make_mfg_csr(dst_nodes, colptr, row, max_num_neighbors,
                                     num_src_nodes)
 
-            return SampledCSC(colptr, row, max_num_neighbors, num_src_nodes)
+            return CSC(colptr, row, num_src_nodes,
+                       dst_max_in_degree=max_num_neighbors)
 
         if LEGACY_MODE:
             return make_fg_csr(colptr, row)
 
-        return StaticCSC(colptr, row)
+        return CSC(colptr, row, num_src_nodes=num_src_nodes)
 
     def get_typed_cugraph(
         self,
@@ -135,15 +130,16 @@ class CuGraphModule(torch.nn.Module):  # pragma: no cover
                                        out_node_types=None, in_node_types=None,
                                        edge_types=edge_type)
 
-            return SampledHeteroCSC(colptr, row, edge_type, max_num_neighbors,
-                                    num_src_nodes, num_edge_types)
+            return HeteroCSC(colptr, row, edge_type, num_src_nodes,
+                             num_edge_types,
+                             dst_max_in_degree=max_num_neighbors)
 
         if LEGACY_MODE:
             return make_fg_csr_hg(colptr, row, n_node_types=0,
                                   n_edge_types=num_edge_types, node_types=None,
                                   edge_types=edge_type)
 
-        return StaticHeteroCSC(colptr, row, edge_type, num_edge_types)
+        return HeteroCSC(colptr, row, edge_type, num_src_nodes, num_edge_types)
 
     def forward(
         self,

torch_geometric/nn/conv/edge_conv.py

@@ -3,13 +3,14 @@ from typing import Callable, Optional, Union
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.inits import reset
 from torch_geometric.typing import Adj, OptTensor, PairOptTensor, PairTensor
 
-try:
+if torch_geometric.typing.WITH_TORCH_CLUSTER:
     from torch_cluster import knn
-except ImportError:
+else:
     knn = None