Repository: pyg-team/pytorch_geometric
Pull request #7481 (2023-05-31 23:46:31)

Make `TGN` use `torch` native when `torch-scatter` not available
in my tested torch native and torch scatter have similar perf for the tgn training example (which includes both forward and back)

```
mine:
root@212f80de0614:/workspace# python3 examples/tgn.py 
Epoch: 01, Loss: 1.1555
Val AP: 0.8242, Val AUC: 0.8313
Test AP: 0.8127, Test AUC: 0.8208
epochtime= 13.63751745223999
Epoch: 02, Loss: 0.9218
Val AP: 0.8302, Val AUC: 0.8499
Test AP: 0.8428, Test AUC: 0.8493
epochtime= 13.39035940170288
Epoch: 03, Loss: 0.8480
Val AP: 0.8654, Val AUC: 0.8696
Test AP: 0.8656, Test AUC: 0.8644
epochtime= 12.575356721878052
Epoch: 04, Loss: 0.8008
Val AP: 0.8784, Val AUC: 0.8783
Test AP: 0.8670, Test AUC: 0.8672
epochtime= 13.801454782485962


original:
Epoch: 01, Loss: 1.0855
Val AP: 0.8646, Val AUC: 0.8787
Test AP: 0.8323, Test AUC: 0.8527
epochtime= 13.7134690284729
Epoch: 02, Loss: 0.8287
Val AP: 0.9079, Val AUC: 0.9113
Test AP: 0.9035, Test AUC: 0.9053
epochtime= 13.247861862182617
Epoch: 03, Loss: 0.7423
Val AP: 0.9243, Val AUC: 0.9269
Test AP: 0.9217, Test AUC: 0.9235
epochtime= 12.532079458236694
Epoch: 04, Loss: 0.6858
Val AP: 0.9380, Val AUC: 0.9351
Test AP: 0.9280, Test AUC: 0.9248
epochtime= 12.426196813583374
```

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7481?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7481](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7481?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (a54f2e9) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/db460561919ca384ef6b10f0ed3b7a8a6dab95e2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (db46056) will **decrease** coverage by `0.03%`.
> The diff coverage is `42.85%`.

```diff
@@            Coverage Diff             @@
##           master    #7481      +/-   ##
==========================================
- Coverage   91.47%   91.44%   -0.03%     
==========================================
  Files         446      446              
  Lines       24778    24789      +11     
==========================================
+ Hits        22666    22669       +3     
- Misses       2112     2120       +8     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7481?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/models/tgn.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7481?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy90Z24ucHk=) | `92.77% <33.33%> (-4.29%)` | :arrow_down: |
| [torch\_geometric/utils/cross\_entropy.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7481?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2Nyb3NzX2VudHJvcHkucHk=) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/torch_geometric/nn/models/tgn.py b/torch_geometric/nn/models/tgn.py
index d224f9b06..24c391fba 100644
--- a/torch_geometric/nn/models/tgn.py
+++ b/torch_geometric/nn/models/tgn.py
@@ -7,6 +7,7 @@ from torch.nn import GRUCell, Linear
 
 from torch_geometric.nn.inits import zeros
 from torch_geometric.utils import scatter
+from torch_geometric.utils.scatter import scatter_argmax
 
 TGNMessageStoreType = Dict[int, Tuple[Tensor, Tensor, Tensor, Tensor]]
 
@@ -194,8 +195,7 @@ class IdentityMessage(torch.nn.Module):
 
 class LastAggregator(torch.nn.Module):
     def forward(self, msg: Tensor, index: Tensor, t: Tensor, dim_size: int):
-        from torch_scatter import scatter_max
-        _, argmax = scatter_max(t, index, dim=0, dim_size=dim_size)
+        argmax = scatter_argmax(t, index, dim=0, dim_size=dim_size)
         out = msg.new_zeros((dim_size, msg.size(-1)))
         mask = argmax < msg.size(0)  # Filter items with at least one entry.
         out[mask] = msg[argmax[mask]]
diff --git a/torch_geometric/utils/scatter.py b/torch_geometric/utils/scatter.py
index 472d19599..5d14ee871 100644
--- a/torch_geometric/utils/scatter.py
+++ b/torch_geometric/utils/scatter.py
@@ -164,3 +164,28 @@ else:  # pragma: no cover
 def broadcast(src: Tensor, ref: Tensor, dim: int) -> Tensor:
     size = ((1, ) * dim) + (-1, ) + ((1, ) * (ref.dim() - dim - 1))
     return src.view(size).expand_as(ref)
+
+
+def scatter_argmax(src: Tensor, index: Tensor, dim: int = 0,
+                   dim_size: Optional[int] = None) -> Tensor:
+
+    if torch_geometric.typing.WITH_TORCH_SCATTER:
+        out = torch_scatter.scatter_max(src, index, dim=dim, dim_size=dim_size)
+        return out[1]
+
+    # Only implemented under certain conditions for now :(
+    assert dim == 0
+    assert src.dim() == 1 and index.dim() == 1
+
+    if dim_size is None:
+        dim_size = index.max() + 1 if index.numel() > 0 else 0
+
+    res = src.new_empty(dim_size)
+    res.scatter_reduce_(0, index, src.detach(), reduce='max',
+                        include_self=False)
+
+    out = index.new_full((dim_size, ), fill_value=dim_size - 1)
+    nonzero = (src == res[index]).nonzero().view(-1)
+    out[index[nonzero]] = nonzero
+
+    return out


Test patch:
diff --git a/test/nn/models/test_tgn.py b/test/nn/models/test_tgn.py
index c9437fa74..cb1f5c6dd 100644
--- a/test/nn/models/test_tgn.py
+++ b/test/nn/models/test_tgn.py
@@ -8,10 +8,8 @@ from torch_geometric.nn.models.tgn import (
     LastAggregator,
     LastNeighborLoader,
 )
-from torch_geometric.testing import withPackage
 
 
-@withPackage('torch_scatter')  # TODO Requires `scatter_argmax` for now.
 def test_tgn():
     memory_dim = 16
     time_dim = 16
diff --git a/test/utils/test_scatter.py b/test/utils/test_scatter.py
index 735a5cee1..d93f17000 100644
--- a/test/utils/test_scatter.py
+++ b/test/utils/test_scatter.py
@@ -2,8 +2,9 @@ import pytest
 import torch
 
 from torch_geometric.profile import benchmark
-from torch_geometric.testing import withCUDA, withPackage
+from torch_geometric.testing import disableExtensions, withCUDA, withPackage
 from torch_geometric.utils import scatter
+from torch_geometric.utils.scatter import scatter_argmax
 
 
 def test_scatter_validate():
@@ -70,6 +71,16 @@ def test_scatter_any(device):
             assert float(out[i, j]) in src[2 * i:2 * i + 2, j].tolist()
 
 
+@withCUDA
+@disableExtensions
+def test_scatter_argmax(device):
+    src = torch.arange(5, device=device)
+    index = torch.tensor([2, 2, 0, 0, 3], device=device)
+
+    argmax = scatter_argmax(src, index, dim_size=6)
+    assert argmax.tolist() == [3, 5, 1, 4, 5, 5]
+
+
 if __name__ == '__main__':
     # Insights on GPU:
     # ================