Repository: pyg-team/pytorch_geometric
Pull request #8525 (2023-12-04 20:14:00)

`[EdgeIndex]` Clean-up, support `reduce=amin` and `reduce=amax`


Patch:
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index 9d25a1134..13961986a 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -31,7 +31,9 @@ else:
         torch.int64,
     }
 
-ReduceType = Literal['sum', 'mean', 'min', 'max']
+ReduceType = Literal['sum', 'mean', 'amin', 'amax', 'add', 'min', 'max']
+PYG_REDUCE = {'add': 'sum', 'amin': 'min', 'amax': 'max'}
+TORCH_REDUCE = {'add': 'sum', 'min': 'amin', 'max': 'amax'}
 
 
 class SortOrder(Enum):
@@ -983,6 +985,7 @@ def _torch_sparse_spmm(
     # `torch-sparse` still provides a faster sparse-dense matrix multiplication
     # code path on GPUs (after all these years...):
     assert torch_geometric.typing.WITH_TORCH_SPARSE
+    reduce = PYG_REDUCE[reduce] if reduce in PYG_REDUCE else reduce
 
     # Optional arguments for backpropagation:
     colptr: Optional[Tensor] = None
@@ -1030,8 +1033,7 @@ class _TorchSPMM(torch.autograd.Function):
         transpose: bool = False,
     ) -> Tensor:
 
-        reduce = 'amin' if reduce == 'min' else reduce
-        reduce = 'amax' if reduce == 'max' else reduce
+        reduce = TORCH_REDUCE[reduce] if reduce in TORCH_REDUCE else reduce
 
         value = value.detach() if value is not None else value
         if other.requires_grad:
@@ -1150,9 +1152,11 @@ def _spmm(
         return _torch_sparse_spmm(input, other, value, reduce, transpose)
 
     if value is not None and value.requires_grad:
+        if torch_geometric.typing.WITH_TORCH_SPARSE:
+            return _torch_sparse_spmm(input, other, value, reduce, transpose)
         return _scatter_spmm(input, other, value, reduce, transpose)
 
-    if reduce == 'sum':
+    if reduce == 'sum' or reduce == 'add':
         return _TorchSPMM.apply(input, other, value, 'sum', transpose)
 
     if reduce == 'mean':
diff --git a/torch_geometric/utils/scatter.py b/torch_geometric/utils/scatter.py
index d9207c90a..c74a0b74f 100644
--- a/torch_geometric/utils/scatter.py
+++ b/torch_geometric/utils/scatter.py
@@ -81,7 +81,7 @@ if torch_geometric.typing.WITH_PT112:  # pragma: no cover
 
         # For "min" and "max" reduction, we prefer `scatter_reduce_` on CPU or
         # in case the input does not require gradients:
-        if reduce == 'min' or reduce == 'max':
+        if reduce in ['min', 'max', 'amin', 'amax']:
             if (not torch_geometric.typing.WITH_TORCH_SCATTER
                     or not src.is_cuda or not src.requires_grad):
 
@@ -92,10 +92,11 @@ if torch_geometric.typing.WITH_PT112:  # pragma: no cover
 
                 index = broadcast(index, src, dim)
                 return src.new_zeros(size).scatter_reduce_(
-                    dim, index, src, reduce=f'a{reduce}', include_self=False)
+                    dim, index, src, reduce=f'a{reduce[-3:]}',
+                    include_self=False)
 
             return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                         reduce=reduce)
+                                         reduce=reduce[-3:])
 
         # For "mul" reduction, we prefer `scatter_reduce_` on CPU:
         if reduce == 'mul':
@@ -138,8 +139,8 @@ else:  # pragma: no cover
                 minimal-sized output tensor according to
                 :obj:`index.max() + 1`. (default: :obj:`None`)
             reduce (str, optional): The reduce operation (:obj:`"sum"`,
-                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`).
-                (default: :obj:`"sum"`)
+                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`,
+                :obj:`"any"`). (default: :obj:`"sum"`)
         """
         if reduce == 'any':
             dim = src.dim() + dim if dim < 0 else dim
@@ -154,6 +155,10 @@ else:  # pragma: no cover
 
         if not torch_geometric.typing.WITH_TORCH_SCATTER:
             raise ImportError("'scatter' requires the 'torch-scatter' package")
+
+        if reduce == 'amin' or reduce == 'amax':
+            reduce = reduce[-3:]
+
         return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
                                      reduce=reduce)
 


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index fa44c28e3..7a20e7cb5 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -11,6 +11,7 @@ from torch_geometric.data import EdgeIndex
 from torch_geometric.data.edge_index import (
     SUPPORTED_DTYPES,
     ReduceType,
+    _scatter_spmm,
     _torch_sparse_spmm,
     _TorchSPMM,
 )
@@ -508,13 +509,12 @@ def test_to_sparse_tensor(device):
 def test_torch_sparse_spmm(device, reduce, transpose):
     adj = EdgeIndex([[0, 1, 1, 2], [2, 0, 1, 2]], device=device)
     adj = adj.sort_by('col' if transpose else 'row').values
-    i = 0 if transpose else 1
 
     # Basic:
     x = torch.randn(3, 1, device=device)
 
     out = _torch_sparse_spmm(adj, x, None, reduce, transpose)
-    exp = scatter(x[adj[i]], adj[1 - i], reduce=reduce)
+    exp = _scatter_spmm(adj, x, None, reduce, transpose)
     assert out.allclose(exp)
 
     # With non-zero values:
@@ -522,7 +522,7 @@ def test_torch_sparse_spmm(device, reduce, transpose):
     value = torch.rand(adj.size(1), device=device)
 
     out = _torch_sparse_spmm(adj, x, value, reduce, transpose)
-    exp = scatter(x[adj[i]] * value.view(-1, 1), adj[1 - i], reduce=reduce)
+    exp = _scatter_spmm(adj, x, value, reduce, transpose)
     assert out.allclose(exp)
 
     # Gradients w.r.t. other:
@@ -532,7 +532,7 @@ def test_torch_sparse_spmm(device, reduce, transpose):
 
     out = _torch_sparse_spmm(adj, x1, None, reduce, transpose)
     out.backward(grad)
-    exp = scatter(x2[adj[i]], adj[1 - i], reduce=reduce)
+    exp = _scatter_spmm(adj, x2, None, reduce, transpose)
     exp.backward(grad)
     assert x1.grad.allclose(x2.grad)
 
@@ -544,7 +544,7 @@ def test_torch_sparse_spmm(device, reduce, transpose):
 
     out = _torch_sparse_spmm(adj, x, value1, reduce, transpose)
     out.backward(grad)
-    exp = scatter(x[adj[i]] * value2.view(-1, 1), adj[1 - i], reduce=reduce)
+    exp = _scatter_spmm(adj, x, value2, reduce, transpose)
     exp.backward(grad)
     assert value1.grad.allclose(value2.grad)
 
@@ -555,14 +555,14 @@ def test_torch_sparse_spmm(device, reduce, transpose):
 def test_torch_spmm(device, reduce, transpose):
     adj = EdgeIndex([[0, 1, 1, 2], [2, 0, 1, 2]], device=device)
     adj, perm = adj.sort_by('col' if transpose else 'row')
-    i = 0 if transpose else 1
 
     # Basic:
     x = torch.randn(3, 2, device=device)
 
-    if (not x.is_cuda and torch_geometric.typing.WITH_PT20) or reduce == 'sum':
+    if ((not x.is_cuda and torch_geometric.typing.WITH_PT20)
+            or reduce in ['sum', 'add']):
         out = _TorchSPMM.apply(adj, x, None, reduce, transpose)
-        exp = scatter(x[adj[i]], adj[1 - i], reduce=reduce)
+        exp = _scatter_spmm(adj, x, None, reduce, transpose)
         assert out.allclose(exp)
     else:
         with pytest.raises(AssertionError):
@@ -572,9 +572,10 @@ def test_torch_spmm(device, reduce, transpose):
     x = torch.randn(3, 1, device=device)
     value = torch.rand(adj.size(1), device=device)
 
-    if (not x.is_cuda and torch_geometric.typing.WITH_PT20) or reduce == 'sum':
+    if ((not x.is_cuda and torch_geometric.typing.WITH_PT20)
+            or reduce in ['sum', 'add']):
         out = _TorchSPMM.apply(adj, x, value, reduce, transpose)
-        exp = scatter(x[adj[i]] * value.view(-1, 1), adj[1 - i], reduce=reduce)
+        exp = _scatter_spmm(adj, x, value, reduce, transpose)
         assert out.allclose(exp)
     else:
         with pytest.raises(AssertionError):
@@ -585,10 +586,10 @@ def test_torch_spmm(device, reduce, transpose):
     x2 = x1.detach().requires_grad_()
     grad = torch.randn_like(x1)
 
-    if reduce == 'sum':
+    if reduce in ['sum', 'add']:
         out = _TorchSPMM.apply(adj, x1, None, reduce, transpose)
         out.backward(grad)
-        exp = scatter(x2[adj[i]], adj[1 - i], reduce=reduce)
+        exp = _scatter_spmm(adj, x2, None, reduce, transpose)
         exp.backward(grad)
         assert x1.grad.allclose(x2.grad)
     else:
@@ -829,13 +830,12 @@ if __name__ == '__main__':
         return adj.matmul(x, reduce=reduce)
 
     def scatter_mm(edge_index, x, reduce):
-        return scatter(x[edge_index[1]], edge_index[0], dim_size=x.size(0),
-                       reduce=reduce)
+        return _scatter_spmm(edge_index, x, reduce=reduce)
 
     funcs = [edge_index_mm, torch_sparse_mm, sparse_tensor_mm, scatter_mm]
     func_names = ['edge_index', 'torch.sparse', 'SparseTensor', 'scatter']
 
-    for reduce in ReduceType.__args__:
+    for reduce in ['sum', 'mean', 'amin', 'amax']:
         func_args = [(edge_index, x, reduce), (adj1, x), (adj2, x, reduce),
                      (edge_index, x, reduce)]
         print(f"reduce='{reduce}':")