Repository: pyg-team/pytorch_geometric
Pull request #7477 (2023-05-31 16:15:23)

Making `RevGNN` example work w/o `torch-sparse`

rusty1s: Thanks for the update. Would it make sense to just use the `ToSparseTensor` transform in this case?

Patch:
diff --git a/.github/workflows/building_rusty1s_conda.yml b/.github/workflows/building_rusty1s_conda.yml
index 296a402cb..d0b88eda2 100644
--- a/.github/workflows/building_rusty1s_conda.yml
+++ b/.github/workflows/building_rusty1s_conda.yml
@@ -11,8 +11,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-18.04, macos-10.15, windows-2019]
-        # We have troube building for Python 3.11 due to version conflicts.
-        python-version: ['3.7', '3.8', '3.9', '3.10']  # '3.11'
+        python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
         torch-version: [1.12.0, 1.13.0, 2.0.0]
         cuda-version: ['cpu', 'cu102', 'cu113', 'cu116', 'cu117', 'cu118']
         exclude:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f2c1e3f16..c7f2f300b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `MovieLens-1M` heterogeneous dataset ([#7479](https://github.com/pyg-team/pytorch_geometric/pull/7479))
+- Added a CPU-based and GPU-based `map_index` implementation ([#7493](https://github.com/pyg-team/pytorch_geometric/pull/7493))
 - Added the `AmazonBook` heterogeneous dataset ([#7483](https://github.com/pyg-team/pytorch_geometric/pull/7483))
 - Added hierarichial heterogeneous GraphSAGE example on OGB-MAG ([#7425](https://github.com/pyg-team/pytorch_geometric/pull/7425))
 - Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452))
diff --git a/examples/rev_gnn.py b/examples/rev_gnn.py
index aef5fcdbe..56fd5fa3a 100644
--- a/examples/rev_gnn.py
+++ b/examples/rev_gnn.py
@@ -14,7 +14,6 @@ from tqdm import tqdm
 import torch_geometric.transforms as T
 from torch_geometric.loader import RandomNodeLoader
 from torch_geometric.nn import GroupAddRev, SAGEConv
-from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import index_to_mask
 
 
@@ -81,9 +80,11 @@ class RevGNN(torch.nn.Module):
 
 from ogb.nodeproppred import Evaluator, PygNodePropPredDataset  # noqa
 
-transform = T.AddSelfLoops()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+transform = T.Compose([T.ToDevice(device), T.ToSparseTensor()])
 root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'products')
-dataset = PygNodePropPredDataset('ogbn-products', root, transform=transform)
+dataset = PygNodePropPredDataset('ogbn-products', root,
+                                 transform=T.AddSelfLoops())
 evaluator = Evaluator(name='ogbn-products')
 
 data = dataset[0]
@@ -97,7 +98,6 @@ train_loader = RandomNodeLoader(data, num_parts=10, shuffle=True,
 # the full batch graph into your GPU:
 test_loader = RandomNodeLoader(data, num_parts=1, num_workers=5)
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = RevGNN(
     in_channels=dataset.num_features,
     hidden_channels=160,
@@ -117,12 +117,11 @@ def train(epoch):
 
     total_loss = total_examples = 0
     for data in train_loader:
-        data = data.to(device)
         optimizer.zero_grad()
 
         # Memory-efficient aggregations:
-        adj_t = SparseTensor.from_edge_index(data.edge_index).t()
-        out = model(data.x, adj_t)[data.train_mask]
+        data = transform(data)
+        out = model(data.x, data.adj_t)[data.train_mask]
         loss = F.cross_entropy(out, data.y[data.train_mask].view(-1))
         loss.backward()
         optimizer.step()
@@ -147,11 +146,9 @@ def test(epoch):
     pbar.set_description(f'Evaluating epoch: {epoch:03d}')
 
     for data in test_loader:
-        data = data.to(device)
-
         # Memory-efficient aggregations
-        adj_t = SparseTensor.from_edge_index(data.edge_index).t()
-        out = model(data.x, adj_t).argmax(dim=-1, keepdim=True)
+        data = transform(data)
+        out = model(data.x, data.adj_t).argmax(dim=-1, keepdim=True)
 
         for split in ['train', 'valid', 'test']:
             mask = data[f'{split}_mask']
diff --git a/torch_geometric/datasets/__init__.py b/torch_geometric/datasets/__init__.py
index 7bf92d16c..1891f59af 100644
--- a/torch_geometric/datasets/__init__.py
+++ b/torch_geometric/datasets/__init__.py
@@ -78,6 +78,7 @@ from .ogb_mag import OGB_MAG
 from .dblp import DBLP
 from .movie_lens import MovieLens
 from .movie_lens_100k import MovieLens100K
+from .movie_lens_1m import MovieLens1M
 from .imdb import IMDB
 from .last_fm import LastFM
 from .hgb_dataset import HGBDataset
@@ -180,6 +181,7 @@ hetero_datasets = [
     'DBLP',
     'MovieLens',
     'MovieLens100K',
+    'MovieLens1M',
     'IMDB',
     'LastFM',
     'HGBDataset',
diff --git a/torch_geometric/datasets/movie_lens_100k.py b/torch_geometric/datasets/movie_lens_100k.py
index 7b29e840a..b7b8c6001 100644
--- a/torch_geometric/datasets/movie_lens_100k.py
+++ b/torch_geometric/datasets/movie_lens_100k.py
@@ -66,7 +66,6 @@ class MovieLens100K(InMemoryDataset):
           - 1
           - 1
     """
-
     url = 'https://files.grouplens.org/datasets/movielens/ml-100k.zip'
 
     def __init__(
@@ -148,7 +147,7 @@ class MovieLens100K(InMemoryDataset):
         edge_index = torch.tensor([src, dst])
         data['user', 'rates', 'movie'].edge_index = edge_index
 
-        rating = torch.from_numpy(df['rating'].values).to(torch.float)
+        rating = torch.from_numpy(df['rating'].values).to(torch.long)
         data['user', 'rates', 'movie'].rating = rating
 
         time = torch.from_numpy(df['timestamp'].values)
diff --git a/torch_geometric/datasets/movie_lens_1m.py b/torch_geometric/datasets/movie_lens_1m.py
new file mode 100644
index 000000000..092c0ff39
--- /dev/null
+++ b/torch_geometric/datasets/movie_lens_1m.py
@@ -0,0 +1,165 @@
+import os
+import os.path as osp
+import shutil
+from typing import Callable, List, Optional
+
+import torch
+
+from torch_geometric.data import (
+    HeteroData,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+
+MOVIE_HEADERS = ["movieId", "title", "genres"]
+USER_HEADERS = ["userId", "gender", "age", "occupation", "zipCode"]
+RATING_HEADERS = ['userId', 'movieId', 'rating', 'timestamp']
+
+
+class MovieLens1M(InMemoryDataset):
+    r"""The MovieLens 1M heterogeneous rating dataset, assembled by GroupLens
+    Research from the `MovieLens web site <https://movielens.org>`__,
+    consisting of movies (3,883 nodes) and users (6,040 nodes) with
+    approximately 1 million ratings between them.
+    User ratings for movies are available as ground truth labels.
+    Features of users and movies are encoded according to the `"Inductive
+    Matrix Completion Based on Graph Neural Networks"
+    <https://arxiv.org/abs/1904.12058>`__ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            every access. (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+
+    **STATS:**
+
+    .. list-table::
+        :widths: 20 10 10 10
+        :header-rows: 1
+
+        * - Node/Edge Type
+          - #nodes/#edges
+          - #features
+          - #tasks
+        * - Movie
+          - 3,883
+          - 18
+          -
+        * - User
+          - 6,040
+          - 30
+          -
+        * - User-Movie
+          - 1,000,209
+          - 1
+          - 1
+    """
+    url = 'https://files.grouplens.org/datasets/movielens/ml-1m.zip'
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+    ):
+        super().__init__(root, transform, pre_transform)
+        self.load(self.processed_paths[0], data_cls=HeteroData)
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['movies.dat', 'users.dat', 'ratings.dat']
+
+    @property
+    def processed_file_names(self) -> str:
+        return 'data.pt'
+
+    def download(self):
+        path = download_url(self.url, self.root)
+        extract_zip(path, self.root)
+        os.remove(path)
+        folder = osp.join(self.root, 'ml-1m')
+        shutil.rmtree(self.raw_dir)
+        os.rename(folder, self.raw_dir)
+
+    def process(self):
+        import pandas as pd
+
+        data = HeteroData()
+
+        # Process movie data:
+        df = pd.read_csv(
+            self.raw_paths[0],
+            sep='::',
+            header=None,
+            index_col='movieId',
+            names=MOVIE_HEADERS,
+            encoding='ISO-8859-1',
+            engine='python',
+        )
+        movie_mapping = {idx: i for i, idx in enumerate(df.index)}
+
+        genres = df['genres'].str.get_dummies('|').values
+        genres = torch.from_numpy(genres).to(torch.float)
+
+        data['movie'].x = genres
+
+        # Process user data:
+        df = pd.read_csv(
+            self.raw_paths[1],
+            sep='::',
+            header=None,
+            index_col='userId',
+            names=USER_HEADERS,
+            dtype='str',
+            encoding='ISO-8859-1',
+            engine='python',
+        )
+        user_mapping = {idx: i for i, idx in enumerate(df.index)}
+
+        age = df['age'].str.get_dummies().values
+        age = torch.from_numpy(age).to(torch.float)
+
+        gender = df['gender'].str.get_dummies().values
+        gender = torch.from_numpy(gender).to(torch.float)
+
+        occupation = df['occupation'].str.get_dummies().values
+        occupation = torch.from_numpy(occupation).to(torch.float)
+
+        data['user'].x = torch.cat([age, gender, occupation], dim=-1)
+
+        # Process rating data:
+        df = pd.read_csv(
+            self.raw_paths[2],
+            sep='::',
+            header=None,
+            names=RATING_HEADERS,
+            encoding='ISO-8859-1',
+            engine='python',
+        )
+
+        src = [user_mapping[idx] for idx in df['userId']]
+        dst = [movie_mapping[idx] for idx in df['movieId']]
+        edge_index = torch.tensor([src, dst])
+        data['user', 'rates', 'movie'].edge_index = edge_index
+
+        rating = torch.from_numpy(df['rating'].values).to(torch.long)
+        data['user', 'rates', 'movie'].rating = rating
+
+        time = torch.from_numpy(df['timestamp'].values)
+        data['user', 'rates', 'movie'].time = time
+
+        data['movie', 'rated_by', 'user'].edge_index = edge_index.flip([0])
+        data['movie', 'rated_by', 'user'].rating = rating
+        data['movie', 'rated_by', 'user'].time = time
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
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/map.py b/torch_geometric/utils/map.py
index cea56b0ed..39db6520d 100644
--- a/torch_geometric/utils/map.py
+++ b/torch_geometric/utils/map.py
@@ -1,16 +1,29 @@
-from typing import Tuple
+import warnings
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
+from torch.utils.dlpack import from_dlpack
 
 
-def map_index(src: Tensor, index: Tensor) -> Tuple[Tensor, Tensor]:
+def map_index(
+    src: Tensor,
+    index: Tensor,
+    max_index: Optional[int] = None,
+    inclusive: bool = False,
+) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Maps indices in :obj:`src` to the positional value of their
     corresponding occurence in :obj:`index`.
+    Indices must be strictly positive.
 
     Args:
         src (torch.Tensor): The source tensor to map.
         index (torch.Tensor): The index tensor that denotes the new mapping.
+        max_index (int, optional): The maximum index value.
+            (default :obj:`None`)
+        inclusive (bool, optional): If set to :obj:`True`, it is assumed that
+            every entry in :obj:`src` has a valid entry in :obj:`index`.
+            Can speed-up computation. (default: :obj:`False`)
 
     :rtype: (:class:`torch.Tensor`, :class:`torch.BoolTensor`)
 
@@ -28,23 +41,106 @@ def map_index(src: Tensor, index: Tensor) -> Tuple[Tensor, Tensor]:
         >>> map_index(src, index)
         (tensor([1, 2, 2, 0]), tensor([True, True, False, True, True]))
     """
-    import pandas as pd
-
-    assert src.dim() == 1 and index.dim() == 1
-    assert not src.is_floating_point()
-    assert not index.is_floating_point()
-
-    arange = pd.RangeIndex(0, index.size(0))
-    df = pd.DataFrame(index=index.detach().cpu().numpy(), data={'out': arange})
-    ser = pd.Series(src.detach().cpu(), name='key')
-    result = df.merge(ser, how='right', left_index=True, right_on='key')
-    out = torch.from_numpy(result['out'].values).to(index.device)
-
-    if out.is_floating_point():
-        mask = torch.isnan(out).logical_not_()
-        out = out[mask].to(index.dtype)
-        return out, mask
-
-    out = out.to(index.dtype)
-    mask = torch.ones_like(out, dtype=torch.bool)
-    return out, mask
+    if src.is_floating_point():
+        raise ValueError(f"Expected 'src' to be an index (got '{src.dtype}')")
+    if index.is_floating_point():
+        raise ValueError(f"Expected 'index' to be an index (got "
+                         f"'{index.dtype}')")
+    if src.device != index.device:
+        raise ValueError(f"Both 'src' and 'index' must be on the same device "
+                         f"(got '{src.device}' and '{index.device}')")
+
+    if max_index is None:
+        max_index = max(src.max(), index.max())
+
+    # If the `max_index` is in a reasonable range, we can accelerate this
+    # operation by creating a helper vector to perform the mapping.
+    # NOTE This will potentially consumes a large chunk of memory
+    # (max_index=10 million => ~75MB), so we cap it at a reasonable size:
+    THRESHOLD = 40_000_000 if src.is_cuda else 10_000_000
+    if max_index <= THRESHOLD:
+        if inclusive:
+            assoc = src.new_empty(max_index + 1)
+        else:
+            assoc = src.new_full((max_index + 1, ), -1)
+        assoc[index] = torch.arange(index.numel(), dtype=src.dtype,
+                                    device=src.device)
+        out = assoc[src]
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask
+
+    WITH_CUDF = False
+    if src.is_cuda:
+        try:
+            import cudf
+            WITH_CUDF = True
+        except ImportError:
+            import pandas as pd
+            warnings.warn("Using CPU-based processing within 'map_index' "
+                          "which may cause slowdowns and device "
+                          "synchronization. Consider installing 'cudf' to "
+                          "accelerate computation")
+    else:
+        import pandas as pd
+
+    if not WITH_CUDF:
+        left_ser = pd.Series(src.cpu().numpy(), name='left_ser')
+        right_ser = pd.Series(
+            index=index.cpu().numpy(),
+            data=pd.RangeIndex(0, index.size(0)),
+            name='right_ser',
+        )
+
+        result = pd.merge(left_ser, right_ser, how='left', left_on='left_ser',
+                          right_index=True)
+
+        out = torch.from_numpy(result['right_ser'].values).to(index.device)
+
+        if out.is_floating_point() and inclusive:
+            raise ValueError("Found invalid entries in 'src' that do not have "
+                             "a corresponding entry in 'index'. Set "
+                             "`inclusive=False` to ignore these entries.")
+
+        if out.is_floating_point():
+            mask = torch.isnan(out).logical_not_()
+            out = out[mask].to(index.dtype)
+            return out, mask
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask
+
+    else:
+        left_ser = cudf.Series(src, name='left_ser')
+        right_ser = cudf.Series(
+            index=index,
+            data=cudf.RangeIndex(0, index.size(0)),
+            name='right_ser',
+        )
+
+        result = cudf.merge(left_ser, right_ser, how='left',
+                            left_on='left_ser', right_index=True, sort=True)
+
+        if inclusive:
+            try:
+                out = from_dlpack(result['right_ser'].to_dlpack())
+            except ValueError:
+                raise ValueError("Found invalid entries in 'src' that do not "
+                                 "have a corresponding entry in 'index'. Set "
+                                 "`inclusive=False` to ignore these entries.")
+        else:
+            out = from_dlpack(result['right_ser'].fillna(-1).to_dlpack())
+
+        out = out[src.argsort().argsort()]  # Restore original order.
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], 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_map.py b/test/utils/test_map.py
index 923bb5b81..df6eb6e54 100644
--- a/test/utils/test_map.py
+++ b/test/utils/test_map.py
@@ -1,21 +1,77 @@
+import pytest
 import torch
 
+from torch_geometric.profile import benchmark
+from torch_geometric.testing import withCUDA
 from torch_geometric.utils.map import map_index
 
 
-def test_map_index():
-    src = torch.tensor([2, 0, 1, 0, 3])
-    index = torch.tensor([3, 2, 0, 1])
+@withCUDA
+@pytest.mark.parametrize('max_index', [3, 100_000_000])
+def test_map_index(device, max_index):
+    src = torch.tensor([2, 0, 1, 0, max_index], device=device)
+    index = torch.tensor([max_index, 2, 0, 1], device=device)
 
-    out, mask = map_index(src, index)
+    out, mask = map_index(src, index, inclusive=True)
+    assert out.device == device
+    assert mask is None
     assert out.tolist() == [1, 2, 3, 2, 0]
-    assert mask.tolist() == [True, True, True, True, True]
 
 
-def test_map_index_na():
-    src = torch.tensor([2, 0, 1, 0, 3])
-    index = torch.tensor([3, 2, 0])
+@withCUDA
+@pytest.mark.parametrize('max_index', [3, 100_000_000])
+def test_map_index_na(device, max_index):
+    src = torch.tensor([2, 0, 1, 0, max_index], device=device)
+    index = torch.tensor([max_index, 2, 0], device=device)
 
-    out, mask = map_index(src, index)
+    out, mask = map_index(src, index, inclusive=False)
+    assert out.device == device
+    assert mask.device == device
     assert out.tolist() == [1, 2, 2, 0]
     assert mask.tolist() == [True, True, False, True, True]
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cuda')
+    args = parser.parse_args()
+
+    src = torch.randint(0, 100_000_000, (100_000, ), device=args.device)
+    index = src.unique()
+
+    def trivial_map(src, index, max_index, inclusive):
+        if max_index is None:
+            max_index = max(src.max(), index.max())
+
+        if inclusive:
+            assoc = src.new_empty(max_index + 1)
+        else:
+            assoc = src.new_full((max_index + 1, ), -1)
+        assoc[index] = torch.arange(index.numel(), device=index.device)
+        out = assoc[src]
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask
+
+    print('Inclusive:')
+    benchmark(
+        funcs=[trivial_map, map_index],
+        func_names=['trivial', 'map_index'],
+        args=(src, index, None, True),
+        num_steps=100,
+        num_warmups=50,
+    )
+
+    print('Exclusive:')
+    benchmark(
+        funcs=[trivial_map, map_index],
+        func_names=['trivial', 'map_index'],
+        args=(src, index[:50_000], None, False),
+        num_steps=100,
+        num_warmups=50,
+    )
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:
     # ================