Repository: pyg-team/pytorch_geometric
Pull request #8753 (2024-01-12 11:52:07)

Do not cache `_alpha` in `GATv2Conv`
Fixes https://github.com/pyg-team/pytorch_geometric/issues/8741

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8753?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`59e9a4e`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/59e9a4e11e2b5ff57bb75282da8e1260357c59f0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.56% compared to head [(`ea5310f`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8753?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.55%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #8753      +/-   ##
==========================================
- Coverage   89.56%   89.55%   -0.01%     
==========================================
  Files         482      482              
  Lines       31113    31110       -3     
==========================================
- Hits        27865    27862       -3     
  Misses       3248     3248              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8753?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ba9ff52eb..e65b8ea63 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Return an empty tensor from `utils.group_argsort` if its input tensor is empty ([#8752](https://github.com/pyg-team/pytorch_geometric/pull/8752))
 - GNN layers are now jittable by default ([#8745](https://github.com/pyg-team/pytorch_geometric/pull/8745))
 - Sparse node features in `NELL` and `AttributedGraphDataset` are now represented as `torch.sparse_csr_tensor` instead of `torch_sparse.SparseTensor` ([#8679](https://github.com/pyg-team/pytorch_geometric/pull/8679))
 - Accelerated mini-batching of `torch.sparse` tensors ([#8670](https://github.com/pyg-team/pytorch_geometric/pull/8670))
diff --git a/torch_geometric/nn/conv/gatv2_conv.py b/torch_geometric/nn/conv/gatv2_conv.py
index ac703c253..1a2a3a8df 100644
--- a/torch_geometric/nn/conv/gatv2_conv.py
+++ b/torch_geometric/nn/conv/gatv2_conv.py
@@ -129,8 +129,6 @@ class GATv2Conv(MessagePassing):
           or :math:`((|\mathcal{V_t}|, H * F_{out}), ((2, |\mathcal{E}|),
           (|\mathcal{E}|, H)))` if bipartite
     """
-    _alpha: OptTensor
-
     def __init__(
         self,
         in_channels: Union[int, Tuple[int, int]],
@@ -191,8 +189,6 @@ class GATv2Conv(MessagePassing):
         else:
             self.register_parameter('bias', None)
 
-        self._alpha = None
-
         self.reset_parameters()
 
     def reset_parameters(self):
@@ -298,12 +294,12 @@ class GATv2Conv(MessagePassing):
                         "simultaneously is currently not yet supported for "
                         "'edge_index' in a 'SparseTensor' form")
 
-        # propagate_type: (x: PairTensor, edge_attr: OptTensor)
-        out = self.propagate(edge_index, x=(x_l, x_r), edge_attr=edge_attr)
+        # edge_updater_type: (x: PairTensor, edge_attr: OptTensor)
+        alpha = self.edge_updater(edge_index, x=(x_l, x_r),
+                                  edge_attr=edge_attr)
 
-        alpha = self._alpha
-        assert alpha is not None
-        self._alpha = None
+        # propagate_type: (x: PairTensor, alpha: Tensor)
+        out = self.propagate(edge_index, x=(x_l, x_r), alpha=alpha)
 
         if self.concat:
             out = out.view(-1, self.heads * self.out_channels)
@@ -326,9 +322,9 @@ class GATv2Conv(MessagePassing):
         else:
             return out
 
-    def message(self, x_j: Tensor, x_i: Tensor, edge_attr: OptTensor,
-                index: Tensor, ptr: OptTensor,
-                size_i: Optional[int]) -> Tensor:
+    def edge_update(self, x_j: Tensor, x_i: Tensor, edge_attr: OptTensor,
+                    index: Tensor, ptr: OptTensor,
+                    size_i: Optional[int]) -> Tensor:
         x = x_i + x_j
 
         if edge_attr is not None:
@@ -342,8 +338,10 @@ class GATv2Conv(MessagePassing):
         x = F.leaky_relu(x, self.negative_slope)
         alpha = (x * self.att).sum(dim=-1)
         alpha = softmax(alpha, index, ptr, size_i)
-        self._alpha = alpha
         alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+        return alpha
+
+    def message(self, x_j: Tensor, alpha: Tensor) -> Tensor:
         return x_j * alpha.unsqueeze(-1)
 
     def __repr__(self) -> str:
diff --git a/torch_geometric/utils/_scatter.py b/torch_geometric/utils/_scatter.py
index 069654f35..9f738b8b7 100644
--- a/torch_geometric/utils/_scatter.py
+++ b/torch_geometric/utils/_scatter.py
@@ -243,11 +243,20 @@ def group_argsort(
             (default: :obj:`False`)
         stable (bool, optional): Controls the relative order of equivalent
             elements. (default: :obj:`False`)
+
+    Example:
+        >>> src = torch.tensor([0, 1, 5, 4, 3, 2, 6, 7, 8])
+        >>> index = torch.tensor([0, 0, 1, 1, 1, 1, 2, 2, 2])
+        >>> group_argsort(src, index)
+        tensor([0, 1, 3, 2, 1, 0, 0, 1, 2])
     """
     # Only implemented under certain conditions for now :(
     assert src.dim() == 1 and index.dim() == 1
     assert dim == 0 or dim == -1
-    assert src.numel() == index.numel() and src.numel() > 0
+    assert src.numel() == index.numel()
+
+    if src.numel() == 0:
+        return torch.zeros_like(src)
 
     # Normalize `src` to range [0, 1]:
     src = src - src.min()


Test patch:
diff --git a/test/nn/conv/test_gatv2_conv.py b/test/nn/conv/test_gatv2_conv.py
index 8198c3671..9d8894746 100644
--- a/test/nn/conv/test_gatv2_conv.py
+++ b/test/nn/conv/test_gatv2_conv.py
@@ -53,7 +53,6 @@ def test_gatv2_conv():
     assert result[1][0].size() == (2, 7)
     assert result[1][1].size() == (7, 2)
     assert result[1][1].min() >= 0 and result[1][1].max() <= 1
-    assert conv._alpha is None
 
     if torch_geometric.typing.WITH_PT113:
         # PyTorch < 1.13 does not support multi-dimensional CSR values :(
@@ -66,7 +65,6 @@ def test_gatv2_conv():
         result = conv(x1, adj2.t(), return_attention_weights=True)
         assert torch.allclose(result[0], out, atol=1e-6)
         assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 7
-        assert conv._alpha is None
 
     if is_full_test():
 
@@ -88,7 +86,6 @@ def test_gatv2_conv():
         assert result[1][0].size() == (2, 7)
         assert result[1][1].size() == (7, 2)
         assert result[1][1].min() >= 0 and result[1][1].max() <= 1
-        assert conv._alpha is None
 
         if torch_geometric.typing.WITH_TORCH_SPARSE:
 
@@ -109,7 +106,6 @@ def test_gatv2_conv():
             result = jit(x1, adj2.t())
             assert torch.allclose(result[0], out, atol=1e-6)
             assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 7
-            assert conv._alpha is None
 
     # Test bipartite message passing:
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
diff --git a/test/utils/test_scatter.py b/test/utils/test_scatter.py
index 1659ecede..c9a32fa03 100644
--- a/test/utils/test_scatter.py
+++ b/test/utils/test_scatter.py
@@ -94,6 +94,10 @@ def test_group_argsort(num_groups, descending, device):
 
     assert torch.equal(out, expected)
 
+    empty_tensor = torch.tensor([], device=device)
+    out = group_argsort(empty_tensor, empty_tensor)
+    assert out.numel() == 0
+
 
 @withCUDA
 def test_scatter_argmax(device):