Repository: pyg-team/pytorch_geometric
Pull request #8448 (2023-11-27 09:50:53)

Add support for `torch.flip` in `EdgeIndex`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8448?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `3 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`8025bc4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/8025bc402abccb0b2fb695cdb32a4cb248d680ab?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.64% compared to head [(`139b18b`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8448?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.64%.

> :exclamation: Current head 139b18b differs from pull request most recent head 7289cb8. Consider uploading reports for the commit 7289cb8 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8448?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/data/edge\_index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8448?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 86.36% | [3 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8448?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    #8448      +/-   ##
==========================================
- Coverage   87.64%   87.64%   -0.01%     
==========================================
  Files         477      477              
  Lines       29163    29185      +22     
==========================================
+ Hits        25559    25578      +19     
- Misses       3604     3607       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8448?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/data/edge_index.py b/torch_geometric/data/edge_index.py
index b6e733e49..2744c717f 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -216,7 +216,7 @@ def apply_(
     out._sparse_size = out._sparse_size
     out._sort_order = out._sort_order
 
-    # Convert metadata:
+    # Convert cache:
     if tensor._rowptr is not None:
         out._rowptr = fn(tensor._rowptr, *args, **kwargs)
     if tensor._colptr is not None:
@@ -280,3 +280,35 @@ def cat(
     out._sparse_size = (num_rows, num_cols)
 
     return out
+
+
+@implements(torch.flip)
+@implements(Tensor.flip)
+def flip(
+    input: EdgeIndex,
+    dims: Union[int, List[int], Tuple[int, ...]],
+) -> Union[EdgeIndex, Tensor]:
+
+    if isinstance(dims, int):
+        dims = [dims]
+    assert isinstance(dims, (tuple, list))
+
+    out = Tensor.__torch_function__(torch.flip, (Tensor, ), (input, dims))
+    out = out.as_subclass(EdgeIndex)
+
+    # Flip metadata and cache:
+    if 0 in dims or -2 in dims:
+        out._sparse_size = input._sparse_size[::-1]
+
+    if len(dims) == 1 and (dims[0] == 0 or dims[0] == -2):
+        if input._sort_order == SortOrder.ROW:
+            out._sort_order = SortOrder.COL
+        elif input._sort_order == SortOrder.COL:
+            out._sort_order = SortOrder.ROW
+
+        out._rowptr = input._colptr
+        out._colptr = input._rowptr
+        out._csr2csc = input._csc2csr
+        out._csc2csr = input._csr2csc
+
+    return out


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 8f1e452af..e2dab1434 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -72,3 +72,22 @@ def test_edge_index_cat():
     out = torch.cat([adj1, adj2], dim=0)
     assert out.size() == (4, 4)
     assert not isinstance(out, EdgeIndex)
+
+
+def test_edge_index_flip():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+    adj.fill_cache()
+
+    out = adj.flip(0)
+    assert isinstance(out, EdgeIndex)
+    assert torch.equal(out, torch.tensor([[1, 0, 2, 1], [0, 1, 1, 2]]))
+    assert out.sparse_size == (None, 3)
+    assert out.sort_order == 'col'
+    assert torch.equal(out._colptr, torch.tensor([0, 1, 3, 4]))
+
+    out = adj.flip([0, 1])
+    assert isinstance(out, EdgeIndex)
+    assert torch.equal(out, torch.tensor([[1, 2, 0, 1], [2, 1, 1, 0]]))
+    assert out.sparse_size == (None, 3)
+    assert out.sort_order is None
+    assert out._colptr is None