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

torch_geometric/nn/conv/eg_conv.py

@@ -81,7 +81,7 @@ class EGConv(MessagePassing):
         self,
         in_channels: int,
         out_channels: int,
-        aggregators: List[str] = ['symnorm'],
+        aggregators: Optional[List[str]] = None,
         num_heads: int = 8,
         num_bases: int = 4,
         cached: bool = False,
@@ -96,23 +96,23 @@ class EGConv(MessagePassing):
                              f"divisible by the number of heads "
                              f"(got {num_heads})")
 
-        for a in aggregators:
-            if a not in ['sum', 'mean', 'symnorm', 'min', 'max', 'var', 'std']:
-                raise ValueError(f"Unsupported aggregator: '{a}'")
-
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.num_heads = num_heads
         self.num_bases = num_bases
         self.cached = cached
         self.add_self_loops = add_self_loops
-        self.aggregators = aggregators
+        self.aggregators = aggregators or ['symnorm']
+
+        for a in self.aggregators:
+            if a not in ['sum', 'mean', 'symnorm', 'min', 'max', 'var', 'std']:
+                raise ValueError(f"Unsupported aggregator: '{a}'")
 
         self.bases_lin = Linear(in_channels,
                                 (out_channels // num_heads) * num_bases,
                                 bias=False, weight_initializer='glorot')
         self.comb_lin = Linear(in_channels,
-                               num_heads * num_bases * len(aggregators))
+                               num_heads * num_bases * len(self.aggregators))
 
         if bias:
             self.bias = Parameter(torch.empty(out_channels))

torch_geometric/nn/conv/gat_conv.py

@@ -107,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`.
 
@@ -137,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')
@@ -151,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:
@@ -176,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)
 
@@ -195,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)
@@ -270,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:
@@ -288,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
@@ -344,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

@@ -110,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`.
 
@@ -141,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)
@@ -154,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):
@@ -181,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)
 
@@ -196,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)
 
@@ -255,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
@@ -267,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)
@@ -305,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/gen_conv.py

@@ -178,7 +178,7 @@ class GENConv(MessagePassing):
             self.lin_dst = Linear(in_channels[1], out_channels, bias=bias)
 
         channels = [out_channels]
-        for i in range(num_layers - 1):
+        for _ in range(num_layers - 1):
             channels.append(out_channels * expansion)
         channels.append(out_channels)
         self.mlp = MLP(channels, norm=norm, bias=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/gravnet_conv.py

@@ -63,7 +63,8 @@ class GravNetConv(MessagePassing):
         if num_workers is not None:
             warnings.warn(
                 "'num_workers' attribute in '{self.__class__.__name__}' is "
-                "deprecated and will be removed in a future release")
+                "deprecated and will be removed in a future release",
+                stacklevel=2)
 
         self.in_channels = in_channels
         self.out_channels = out_channels

torch_geometric/nn/conv/hetero_conv.py

@@ -77,7 +77,8 @@ class HeteroConv(torch.nn.Module):
                 f"There exist node types ({src_node_types - dst_node_types}) "
                 f"whose representations do not get updated during message "
                 f"passing as they do not occur as destination type in any "
-                f"edge type. This may lead to unexpected behavior.")
+                f"edge type. This may lead to unexpected behavior.",
+                stacklevel=2)
 
         self.convs = ModuleDict(convs)
         self.aggr = aggr
@@ -102,7 +103,7 @@ class HeteroConv(torch.nn.Module):
                 individual edge type, either as a :class:`torch.Tensor` of
                 shape :obj:`[2, num_edges]` or a
                 :class:`torch_sparse.SparseTensor`.
-            *args_dict (optional): Additional forward arguments of invididual
+            *args_dict (optional): Additional forward arguments of individual
                 :class:`torch_geometric.nn.conv.MessagePassing` layers.
             **kwargs_dict (optional): Additional forward arguments of
                 individual :class:`torch_geometric.nn.conv.MessagePassing`