Repository: pyg-team/pytorch_geometric
Pull request #8680 (2023-12-28 18:10:28)

Remove `SparseTensor` dependency in `FusedGATConv`

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

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8680?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/utils/sparse.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8680?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NwYXJzZS5weQ==) | 66.66% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8680?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    #8680      +/-   ##
==========================================
- Coverage   89.47%   89.47%   -0.01%     
==========================================
  Files         481      481              
  Lines       30977    30983       +6     
==========================================
+ Hits        27718    27723       +5     
- Misses       3259     3260       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8680?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/nn/conv/fused_gat_conv.py b/torch_geometric/nn/conv/fused_gat_conv.py
index ef2daa69a..5b8ed378d 100644
--- a/torch_geometric/nn/conv/fused_gat_conv.py
+++ b/torch_geometric/nn/conv/fused_gat_conv.py
@@ -4,7 +4,8 @@ import torch
 from torch import Tensor
 
 from torch_geometric.nn.conv import GATConv
-from torch_geometric.typing import SparseTensor
+from torch_geometric.utils import sort_edge_index
+from torch_geometric.utils.sparse import index2ptr
 
 
 class FusedGATConv(GATConv):  # pragma: no cover
@@ -55,17 +56,19 @@ class FusedGATConv(GATConv):  # pragma: no cover
             size ((int, int), optional): The shape of :obj:`edge_index` in each
                 dimension. (default: :obj:`None`)
         """
-        value = torch.arange(edge_index.size(1), dtype=torch.int,
-                             device=edge_index.device)
+        edge_index = edge_index.to(torch.int)
 
-        # TODO Drop `torch-sparse` requirement.
-        adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=size)
-        adj.set_value_(value, layout='csr')
+        edge_index = sort_edge_index(edge_index, sort_by_row=True)
+        rowptr = index2ptr(edge_index[0], size=size[0] if size else None)
+        col = edge_index[1]
 
-        rowptr, col, _ = adj.csr()
-        colptr, row, perm = adj.csc()
+        device = edge_index.device
+        perm = torch.arange(edge_index.size(1), dtype=torch.int, device=device)
+        edge_index, perm = sort_edge_index(edge_index, perm, sort_by_row=False)
+        row = edge_index[0]
+        colptr = index2ptr(edge_index[1], size=size[1] if size else None)
 
-        return (rowptr.int(), col.int()), (row.int(), colptr.int()), perm
+        return (rowptr, col), (row, colptr), perm
 
     def forward(
         self,
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index ad68ce8c3..1dabc0cca 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -488,7 +488,10 @@ def ptr2index(ptr: Tensor) -> Tensor:
     return ind.repeat_interleave(ptr[1:] - ptr[:-1])
 
 
-def index2ptr(index: Tensor, size: int) -> Tensor:
+def index2ptr(index: Tensor, size: Optional[int] = None) -> Tensor:
+    if size is None:
+        size = int(index.max()) + 1 if index.numel() > 0 else 0
+
     return torch._convert_indices_from_coo_to_csr(
         index, size, out_int32=index.dtype == torch.int32)
 


Test patch:
diff --git a/test/nn/conv/test_fused_gat_conv.py b/test/nn/conv/test_fused_gat_conv.py
index 92709751c..71aed067f 100644
--- a/test/nn/conv/test_fused_gat_conv.py
+++ b/test/nn/conv/test_fused_gat_conv.py
@@ -4,9 +4,26 @@ from torch_geometric.nn import FusedGATConv
 from torch_geometric.testing import onlyCUDA, withPackage
 
 
+def test_to_graph_format() -> None:
+    edge_index = torch.tensor([[1, 0, 2, 3], [0, 0, 1, 1]])
+
+    csr, csc, perm = FusedGATConv.to_graph_format(edge_index, size=(4, 4))
+
+    assert csr[0].dtype == torch.int
+    assert torch.equal(csr[0], torch.tensor([0, 1, 2, 3, 4], dtype=torch.int))
+    assert csr[1].dtype == torch.int
+    assert torch.equal(csr[1], torch.tensor([0, 0, 1, 1], dtype=torch.int))
+    assert csc[0].dtype == torch.int
+    assert torch.equal(csc[0], torch.tensor([0, 1, 2, 3], dtype=torch.int))
+    assert csc[1].dtype == torch.int
+    assert torch.equal(csc[1], torch.tensor([0, 2, 4, 4, 4], dtype=torch.int))
+    assert perm.dtype == torch.int
+    assert torch.equal(perm, torch.tensor([0, 1, 2, 3], dtype=torch.int))
+
+
 @onlyCUDA
 @withPackage('dgNN')
-def test_fused_gat_conv():
+def test_fused_gat_conv() -> None:
     device = torch.device('cuda')
 
     x = torch.randn(4, 8, device=device)
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
index 19cb1047b..964619078 100644
--- a/torch_geometric/testing/decorators.py
+++ b/torch_geometric/testing/decorators.py
@@ -159,12 +159,14 @@ def withPackage(*args: str) -> Callable:
     """
     na_packages = set(package for package in args if not has_package(package))
 
+    if len(na_packages) == 1:
+        reason = f"Package {list(na_packages)[0]} not found"
+    else:
+        reason = f"Packages {na_packages} not found"
+
     def decorator(func: Callable) -> Callable:
         import pytest
-        return pytest.mark.skipif(
-            len(na_packages) > 0,
-            reason=f"Package(s) {na_packages} are not installed",
-        )(func)
+        return pytest.mark.skipif(len(na_packages) > 0, reason=reason)(func)
 
     return decorator