Repository: pyg-team/pytorch_geometric
Pull request #8650 (2023-12-20 18:58:26)

Added support for `EdgeIndex.sub`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8650?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`1ae41f4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/1ae41f4fb3e9949b81f511ad5d2e4017fd09f28f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.42% compared to head [(`75832c3`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8650?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.42%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8650?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/edge\_index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8650?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2VkZ2VfaW5kZXgucHk=) | 95.34% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8650?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

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


```diff
@@           Coverage Diff           @@
##           master    #8650   +/-   ##
=======================================
  Coverage   89.42%   89.42%           
=======================================
  Files         481      481           
  Lines       30757    30800   +43     
=======================================
+ Hits        27503    27543   +40     
- Misses       3254     3257    +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8650?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/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index c653304cd..c38fc00cf 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -90,6 +90,21 @@ def maybe_add(
                  for v, o in zip(value, other))
 
 
+def maybe_sub(
+    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 "
@@ -1224,6 +1239,78 @@ def add_(
     return postprocess_add_(input, other, output, alpha)
 
 
+def postprocess_sub_(
+    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_sub(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_sub(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_sub(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
+
+    return output
+
+
+@implements(torch.sub)
+@implements(Tensor.sub)
+def sub(
+    input: EdgeIndex,
+    other: Union[int, Tensor],
+    *,
+    alpha: int = 1,
+    out: Optional[Tensor] = None,
+) -> Union[EdgeIndex, Tensor]:
+
+    output = Tensor.__torch_function__(  #
+        torch.sub, (Tensor, ), (input, other), dict(alpha=alpha, out=out))
+
+    return postprocess_sub_(input, other, output, alpha)
+
+
+@implements(Tensor.sub_)
+def sub_(
+    input: EdgeIndex,
+    other: Union[int, Tensor],
+    *,
+    alpha: int = 1,
+) -> Union[EdgeIndex, Tensor]:
+
+    output = Tensor.__torch_function__(  #
+        Tensor.sub_, (Tensor, ), (input, other), dict(alpha=alpha))
+
+    return postprocess_sub_(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 c35f996af..68ee5a1d4 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -673,6 +673,50 @@ def test_add(dtype, device, is_undirected):
     assert adj.sparse_size() == (5, 5)
 
 
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('is_undirected', IS_UNDIRECTED)
+def test_sub(dtype, device, is_undirected):
+    kwargs = dict(dtype=dtype, device=device, is_undirected=is_undirected)
+    adj = EdgeIndex([[4, 5, 5, 6], [5, 4, 6, 5]], sparse_size=(7, 7), **kwargs)
+
+    out = torch.sub(adj, 2, alpha=2)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[0, 1, 1, 2], [1, 0, 2, 1]], device=device))
+    assert out.is_undirected == is_undirected
+    assert out.sparse_size() == (3, 3)
+
+    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], [4, 3, 5, 4]], device=device))
+    assert not out.is_undirected
+    assert out.sparse_size() == (5, 6)
+
+    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.sub(adj)
+    assert isinstance(out, EdgeIndex)
+    assert out.equal(tensor([[0, 0, 0, 0], [0, 0, 0, 0]], device=device))
+    assert not out.is_undirected
+    assert out.sparse_size() == (None, None)
+
+    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__)