pyg-nightly 2.6.0.dev20240318__py3-none-any.whl → 2.7.0.dev20250115__py3-none-any.whl

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
 
 

torch_geometric/nn/conv/fused_gat_conv.py

@@ -3,9 +3,9 @@ from typing import Optional, Tuple
 import torch
 from torch import Tensor
 
+from torch_geometric.index import index2ptr
 from torch_geometric.nn.conv import GATConv
 from torch_geometric.utils import sort_edge_index
-from torch_geometric.utils.sparse import index2ptr
 
 
 class FusedGATConv(GATConv):  # pragma: no cover

torch_geometric/nn/conv/gat_conv.py

@@ -37,9 +37,8 @@ class GATConv(MessagePassing):
     <https://arxiv.org/abs/1710.10903>`_ paper.
 
     .. math::
-        \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}_{s}\mathbf{x}_{i} +
-        \sum_{j \in \mathcal{N}(i)}
-        \alpha_{i,j}\mathbf{\Theta}_{t}\mathbf{x}_{j},
+        \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i) \cup \{ i \}}
+        \alpha_{i,j}\mathbf{\Theta}_t\mathbf{x}_{j},
 
     where the attention coefficients :math:`\alpha_{i,j}` are computed as
 
@@ -108,6 +107,8 @@ class GATConv(MessagePassing):
             :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`"mean"`)
         bias (bool, optional): If set to :obj:`False`, the layer will not learn
             an additive bias. (default: :obj:`True`)
+        residual (bool, optional): If set to :obj:`True`, the layer will add
+            a learnable skip-connection. (default: :obj:`False`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.conv.MessagePassing`.
 
@@ -138,6 +139,7 @@ class GATConv(MessagePassing):
         edge_dim: Optional[int] = None,
         fill_value: Union[float, Tensor, str] = 'mean',
         bias: bool = True,
+        residual: bool = False,
         **kwargs,
     ):
         kwargs.setdefault('aggr', 'add')
@@ -152,6 +154,7 @@ class GATConv(MessagePassing):
         self.add_self_loops = add_self_loops
         self.edge_dim = edge_dim
         self.fill_value = fill_value
+        self.residual = residual
 
         # In case we are operating in bipartite graphs, we apply separate
         # transformations 'lin_src' and 'lin_dst' to source and target nodes:
@@ -177,10 +180,22 @@ class GATConv(MessagePassing):
             self.lin_edge = None
             self.register_parameter('att_edge', None)
 
-        if bias and concat:
-            self.bias = Parameter(torch.empty(heads * out_channels))
-        elif bias and not concat:
-            self.bias = Parameter(torch.empty(out_channels))
+        # The number of output channels:
+        total_out_channels = out_channels * (heads if concat else 1)
+
+        if residual:
+            self.res = Linear(
+                in_channels
+                if isinstance(in_channels, int) else in_channels[1],
+                total_out_channels,
+                bias=False,
+                weight_initializer='glorot',
+            )
+        else:
+            self.register_parameter('res', None)
+
+        if bias:
+            self.bias = Parameter(torch.empty(total_out_channels))
         else:
             self.register_parameter('bias', None)
 
@@ -196,6 +211,8 @@ class GATConv(MessagePassing):
             self.lin_dst.reset_parameters()
         if self.lin_edge is not None:
             self.lin_edge.reset_parameters()
+        if self.res is not None:
+            self.res.reset_parameters()
         glorot(self.att_src)
         glorot(self.att_dst)
         glorot(self.att_edge)
@@ -271,11 +288,16 @@ class GATConv(MessagePassing):
 
         H, C = self.heads, self.out_channels
 
+        res: Optional[Tensor] = None
+
         # We first transform the input node features. If a tuple is passed, we
         # transform source and target node features via separate weights:
         if isinstance(x, Tensor):
             assert x.dim() == 2, "Static graphs not supported in 'GATConv'"
 
+            if self.res is not None:
+                res = self.res(x)
+
             if self.lin is not None:
                 x_src = x_dst = self.lin(x).view(-1, H, C)
             else:
@@ -289,6 +311,9 @@ class GATConv(MessagePassing):
             x_src, x_dst = x
             assert x_src.dim() == 2, "Static graphs not supported in 'GATConv'"
 
+            if x_dst is not None and self.res is not None:
+                res = self.res(x_dst)
+
             if self.lin is not None:
                 # If the module is initialized as non-bipartite, we expect that
                 # source and destination node features have the same shape and
@@ -345,6 +370,9 @@ class GATConv(MessagePassing):
         else:
             out = out.mean(dim=1)
 
+        if res is not None:
+            out = out + res
+
         if self.bias is not None:
             out = out + self.bias
 

torch_geometric/nn/conv/gatv2_conv.py

@@ -41,8 +41,7 @@ class GATv2Conv(MessagePassing):
     In contrast, in :class:`GATv2`, every node can attend to any other node.
 
     .. math::
-        \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}_{s}\mathbf{x}_{i} +
-        \sum_{j \in \mathcal{N}(i)}
+        \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i) \cup \{ i \}}
         \alpha_{i,j}\mathbf{\Theta}_{t}\mathbf{x}_{j},
 
     where the attention coefficients :math:`\alpha_{i,j}` are computed as
@@ -111,6 +110,8 @@ class GATv2Conv(MessagePassing):
             will be applied to the source and the target node of every edge,
             *i.e.* :math:`\mathbf{\Theta}_{s} = \mathbf{\Theta}_{t}`.
             (default: :obj:`False`)
+        residual (bool, optional): If set to :obj:`True`, the layer will add
+            a learnable skip-connection. (default: :obj:`False`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.conv.MessagePassing`.
 
@@ -142,6 +143,7 @@ class GATv2Conv(MessagePassing):
         fill_value: Union[float, Tensor, str] = 'mean',
         bias: bool = True,
         share_weights: bool = False,
+        residual: bool = False,
         **kwargs,
     ):
         super().__init__(node_dim=0, **kwargs)
@@ -155,6 +157,7 @@ class GATv2Conv(MessagePassing):
         self.add_self_loops = add_self_loops
         self.edge_dim = edge_dim
         self.fill_value = fill_value
+        self.residual = residual
         self.share_weights = share_weights
 
         if isinstance(in_channels, int):
@@ -182,10 +185,22 @@ class GATv2Conv(MessagePassing):
         else:
             self.lin_edge = None
 
-        if bias and concat:
-            self.bias = Parameter(torch.empty(heads * out_channels))
-        elif bias and not concat:
-            self.bias = Parameter(torch.empty(out_channels))
+        # The number of output channels:
+        total_out_channels = out_channels * (heads if concat else 1)
+
+        if residual:
+            self.res = Linear(
+                in_channels
+                if isinstance(in_channels, int) else in_channels[1],
+                total_out_channels,
+                bias=False,
+                weight_initializer='glorot',
+            )
+        else:
+            self.register_parameter('res', None)
+
+        if bias:
+            self.bias = Parameter(torch.empty(total_out_channels))
         else:
             self.register_parameter('bias', None)
 
@@ -197,6 +212,8 @@ class GATv2Conv(MessagePassing):
         self.lin_r.reset_parameters()
         if self.lin_edge is not None:
             self.lin_edge.reset_parameters()
+        if self.res is not None:
+            self.res.reset_parameters()
         glorot(self.att)
         zeros(self.bias)
 
@@ -256,10 +273,16 @@ class GATv2Conv(MessagePassing):
         """
         H, C = self.heads, self.out_channels
 
+        res: Optional[Tensor] = None
+
         x_l: OptTensor = None
         x_r: OptTensor = None
         if isinstance(x, Tensor):
             assert x.dim() == 2
+
+            if self.res is not None:
+                res = self.res(x)
+
             x_l = self.lin_l(x).view(-1, H, C)
             if self.share_weights:
                 x_r = x_l
@@ -268,6 +291,10 @@ class GATv2Conv(MessagePassing):
         else:
             x_l, x_r = x[0], x[1]
             assert x[0].dim() == 2
+
+            if x_r is not None and self.res is not None:
+                res = self.res(x_r)
+
             x_l = self.lin_l(x_l).view(-1, H, C)
             if x_r is not None:
                 x_r = self.lin_r(x_r).view(-1, H, C)
@@ -306,6 +333,9 @@ class GATv2Conv(MessagePassing):
         else:
             out = out.mean(dim=1)
 
+        if res is not None:
+            out = out + res
+
         if self.bias is not None:
             out = out + self.bias
 

torch_geometric/nn/conv/general_conv.py

@@ -70,7 +70,7 @@ class GeneralConv(MessagePassing):
         self,
         in_channels: Union[int, Tuple[int, int]],
         out_channels: Optional[int],
-        in_edge_channels: int = None,
+        in_edge_channels: Optional[int] = None,
         aggr: str = "add",
         skip_linear: str = False,
         directed_msg: bool = True,

torch_geometric/nn/conv/graph_conv.py

@@ -1,7 +1,9 @@
-from typing import Tuple, Union
+from typing import Final, Tuple, Union
 
+import torch
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import Adj, OptPairTensor, OptTensor, Size
@@ -44,6 +46,8 @@ class GraphConv(MessagePassing):
         - **output:** node features :math:`(|\mathcal{V}|, F_{out})` or
           :math:`(|\mathcal{V}_t|, F_{out})` if bipartite
     """
+    SUPPORTS_FUSED_EDGE_INDEX: Final[bool] = True
+
     def __init__(
         self,
         in_channels: Union[int, Tuple[int, int]],
@@ -90,5 +94,19 @@ class GraphConv(MessagePassing):
     def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
         return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j
 
-    def message_and_aggregate(self, adj_t: Adj, x: OptPairTensor) -> Tensor:
-        return spmm(adj_t, x[0], reduce=self.aggr)
+    def message_and_aggregate(
+        self,
+        edge_index: Adj,
+        x: OptPairTensor,
+        edge_weight: OptTensor,
+    ) -> Tensor:
+
+        if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+            return edge_index.matmul(
+                other=x[0],
+                input_value=edge_weight,
+                reduce=self.aggr,
+                transpose=True,
+            )
+
+        return spmm(edge_index, x[0], reduce=self.aggr)

torch_geometric/nn/conv/gravnet_conv.py

@@ -4,14 +4,15 @@ from typing import Optional, Union
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import OptPairTensor  # noqa
 from torch_geometric.typing import OptTensor, PairOptTensor, PairTensor
 
-try:
+if torch_geometric.typing.WITH_TORCH_CLUSTER:
     from torch_cluster import knn
-except ImportError:
+else:
     knn = None