Repository: pyg-team/pytorch_geometric
Pull request #8649 (2023-12-20 18:39:49)

Added support for `EdgeIndex.add`


Patch:
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 44491d8ed..c653304cd 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -8,6 +8,7 @@ from typing import (
     List,
     Literal,
     Optional,
+    Sequence,
     Set,
     Tuple,
     Type,
@@ -62,7 +63,11 @@ def implements(torch_function: Callable) -> Callable:
     return decorator
 
 
-def set_tuple_item(values: Tuple, dim: int, value: Any) -> Tuple:
+def set_tuple_item(
+    values: Tuple[Any, ...],
+    dim: int,
+    value: Any,
+) -> Tuple[Any, ...]:
     if dim < -len(values) or dim >= len(values):
         raise IndexError("tuple index out of range")
 
@@ -70,6 +75,21 @@ def set_tuple_item(values: Tuple, dim: int, value: Any) -> Tuple:
     return values[:dim] + (value, ) + values[dim + 1:]
 
 
+def maybe_add(
+    value: Sequence[Optional[int]],
+    other: Union[int, Sequence[Optional[int]]],
+    alpha: int = 1,
+) -> Tuple[Optional[int], ...]:
+
+    if isinstance(other, int):
+        return tuple(v + alpha * other if v is not None else None
+                     for v in value)
+
+    assert len(value) == len(other)
+    return tuple(v + alpha * o if v is not None and o is not None else None
+                 for v, o in zip(value, other))
+
+
 def assert_valid_dtype(tensor: Tensor) -> None:
     if tensor.dtype not in SUPPORTED_DTYPES:
         raise ValueError(f"'EdgeIndex' holds an unsupported data type "
@@ -1127,6 +1147,83 @@ def getitem(input: EdgeIndex, index: Any) -> Union[EdgeIndex, Tensor]:
     return out
 
 
+def postprocess_add_(
+    input: EdgeIndex,
+    other: Union[int, Tensor],
+    out: Tensor,
+    alpha: int = 1,
+) -> Union[EdgeIndex, Tensor]:
+
+    if out.dtype not in SUPPORTED_DTYPES:
+        return out
+    if out.dim() != 2 or out.size(0) != 2:
+        return out
+
+    output: EdgeIndex = out.as_subclass(EdgeIndex)
+
+    if isinstance(other, int):
+        size = maybe_add(input._sparse_size, other, alpha)
+        assert len(size) == 2
+        output._sparse_size = size
+        output._sort_order = input._sort_order
+        output._is_undirected = input.is_undirected
+        output._T_perm = input._T_perm
+
+    elif isinstance(other, Tensor) and other.numel() <= 1:
+        size = maybe_add(input._sparse_size, int(other), alpha)
+        assert len(size) == 2
+        output._sparse_size = size
+        output._sort_order = input._sort_order
+        output._is_undirected = input.is_undirected
+        output._T_perm = input._T_perm
+
+    elif isinstance(other, Tensor) and other.size() == (2, 1):
+        size = maybe_add(input._sparse_size, other.view(-1).tolist(), alpha)
+        assert len(size) == 2
+        output._sparse_size = size
+        output._sort_order = input._sort_order
+        output._T_perm = input._T_perm
+        if torch.equal(other[0], other[1]):
+            output._is_undirected = input.is_undirected
+
+    elif isinstance(other, EdgeIndex):
+        size = maybe_add(input._sparse_size, other._sparse_size, alpha)
+        assert len(size) == 2
+        output._sparse_size = size
+
+    return output
+
+
+@implements(torch.add)
+@implements(Tensor.add)
+def add(
+    input: EdgeIndex,
+    other: Union[int, Tensor],
+    *,
+    alpha: int = 1,
+    out: Optional[Tensor] = None,
+) -> Union[EdgeIndex, Tensor]:
+
+    output = Tensor.__torch_function__(  #
+        torch.add, (Tensor, ), (input, other), dict(alpha=alpha, out=out))
+
+    return postprocess_add_(input, other, output, alpha)
+
+
+@implements(Tensor.add_)
+def add_(
+    input: EdgeIndex,
+    other: Union[int, Tensor],
+    *,
+    alpha: int = 1,
+) -> Union[EdgeIndex, Tensor]:
+
+    output = Tensor.__torch_function__(  #
+        Tensor.add_, (Tensor, ), (input, other), dict(alpha=alpha))
+
+    return postprocess_add_(input, other, output, alpha)
+
+
 # Sparse-Dense Matrix Multiplication ##########################################
 
 


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index f8844c81d..c35f996af 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -68,7 +68,7 @@ def test_basic(dtype, device):
     assert out.dtype == dtype
     assert out.device == device
 
-    out = adj + 1
+    out = adj * 1
     assert not isinstance(out, EdgeIndex)
     assert out.dtype == dtype
     assert out.device == device
@@ -629,6 +629,50 @@ def test_to_sparse_tensor(device):
     assert col.equal(adj[1])
 
 
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('is_undirected', IS_UNDIRECTED)
+def test_add(dtype, device, is_undirected):
+    kwargs = dict(dtype=dtype, device=device, is_undirected=is_undirected)
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3), **kwargs)
+
+    out = torch.add(adj, 2, alpha=2)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[4, 5, 5, 6], [5, 4, 6, 5]], device=device))
+    assert out.is_undirected == is_undirected
+    assert out.sparse_size() == (7, 7)
+
+    out = adj + torch.tensor([2], dtype=dtype, device=device)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
+    assert out.is_undirected == is_undirected
+    assert out.sparse_size() == (5, 5)
+
+    out = adj + torch.tensor([[2], [1]], dtype=dtype, device=device)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[2, 3, 3, 4], [2, 1, 3, 2]], device=device))
+    assert not out.is_undirected
+    assert out.sparse_size() == (5, 4)
+
+    out = adj + torch.tensor([[2], [2]], dtype=dtype, device=device)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
+    assert out.is_undirected == is_undirected
+    assert out.sparse_size() == (5, 5)
+
+    out = adj.add(adj)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[0, 2, 2, 4], [2, 0, 4, 2]], device=device))
+    assert not out.is_undirected
+    assert out.sparse_size() == (6, 6)
+
+    adj += 2
+    assert isinstance(adj, EdgeIndex)
+    assert adj.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
+    assert adj.is_undirected == is_undirected
+    assert adj.sparse_size() == (5, 5)
+
+
 @withCUDA
 @withPackage('torch_sparse')
 @pytest.mark.parametrize('reduce', ReduceType.__args__)