Repository: pyg-team/pytorch_geometric
Pull request #7089 (2023-03-31 10:24:52)

`SparseTensor` support for `trim_to_layer`
It provides support for SparseTensor data format for the trim_to_layer.py functionality. 
The trimming function is isolated in a separate function, but still contained in trim_to_layer.py, as it is useful to that only.
Tests are included extending the test file used for trim_to_layer.py.
Some tests already present in a previous PR were slightly modified as to simulate a adj_t matrix that originates from a BFS traversal of a graph, starting the BFS traversal from the target node (assuming Batch Size = 1) 

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7089?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7089](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7089?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (854f399) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/707fee2f9d7ff56a4792693eaea0134390749eef?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (707fee2) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head 854f399 differs from pull request most recent head f1e4f7f. Consider uploading reports for the commit f1e4f7f to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7089   +/-   ##
=======================================
  Coverage   91.25%   91.26%           
=======================================
  Files         436      436           
  Lines       23911    23932   +21     
=======================================
+ Hits        21821    21842   +21     
  Misses       2090     2090           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7089?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/\_\_init\_\_.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7089?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/trim\_to\_layer.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7089?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3RyaW1fdG9fbGF5ZXIucHk=) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

andreazanetti: @rusty1s 
with reference to the failing tests, I have some questions as it is not clear to me how we would like to manage the absence of torch_sparse in the requirements in PyG.

To my understanding, that absence generates issues with this contribution in a number of tests:
-docs/readthedocs.org:pytorch-geometric — [Read the Docs build failed!](https://readthedocs.org/projects/pytorch-geometric/builds/19969696/) it fails as torch_sparse is missing
- Testing PyTorch nightly / latest_pytest (pull_request) [Failing after 2m ](https://github.com/pyg-team/pytorch_geometric/actions/runs/4574073196/jobs/8075260403?pr=7089 it fails for the same reason)
- the same holds true for other failures above like Minimal Installation / import (pull_request) and make_html

the only one not impacted by the absence of torch_sparse is the[ changelog failure](https://github.com/pyg-team/pytorch_geometric/actions/runs/4574364341/jobs/8075895143?pr=7089), which was however present also in other PRs recently merged... so I guess that should not be blocking for now.

How would you like me to go on? 
Should I try to make the whole contribution conditional to the presence (=already installed) of torch_sparse in the user's env?
Thank you!!
rusty1s: You need to import from `torch_geometric.typing` for `SparseTensor` dependency, and any tests need to be gated behind `@withPackage('torch_sparse')`. Hope this clarifies your doubts.
andreazanetti: > You need to import from `torch_geometric.typing` for `SparseTensor` dependency, and any tests need to be gated behind `@withPackage('torch_sparse')`. Hope this clarifies your doubts.

Sure, it clarifies, thanks!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 52f8c1726..e57a98bfb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `SparseTensor` support to`trim_to_layer` function ([#7089](https://github.com/pyg-team/pytorch_geometric/pull/7089))
 - Added instructions for ROCm build wheels ([#7143](https://github.com/pyg-team/pytorch_geometric/pull/7143))
 - Added a `ComposeFilters` class to compose `pre_filter` functions in `Dataset` ([#7097](https://github.com/pyg-team/pytorch_geometric/pull/7097))
 - Added a time-step aware variant of the `EllipticBitcoinDataset` called `EllipticBitcoinTemporalDataset` ([#7011](https://github.com/pyg-team/pytorch_geometric/pull/7011))
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 797c58ff7..769dc0218 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -35,7 +35,7 @@ except (ImportError, OSError) as e:
 
 try:
     import torch_sparse  # noqa
-    from torch_sparse import SparseTensor
+    from torch_sparse import SparseStorage, SparseTensor
     WITH_TORCH_SPARSE = True
 except (ImportError, OSError) as e:
     if isinstance(e, OSError):
@@ -43,6 +43,10 @@ except (ImportError, OSError) as e:
                       f"Disabling its usage. Stacktrace: {e}")
     WITH_TORCH_SPARSE = False
 
+    class SparseStorage:
+        def __init__(*args, **kwargs):
+            raise ImportError("'SparseStorage' requires 'torch-sparse'")
+
     class SparseTensor:
         def __init__(
             self,
diff --git a/torch_geometric/utils/trim_to_layer.py b/torch_geometric/utils/trim_to_layer.py
index 058bb40da..dd2ce60b8 100644
--- a/torch_geometric/utils/trim_to_layer.py
+++ b/torch_geometric/utils/trim_to_layer.py
@@ -8,6 +8,8 @@ from torch_geometric.typing import (
     MaybeHeteroEdgeTensor,
     MaybeHeteroNodeTensor,
     NodeType,
+    SparseStorage,
+    SparseTensor,
 )
 
 
@@ -45,6 +47,7 @@ def trim_to_layer(
     if layer <= 0:
         return x, edge_index, edge_attr
 
+    # TODO Support `SparseTensor` for heterogeneous graphs.
     if isinstance(num_sampled_edges_per_hop, dict):
         x = {
             k: v.narrow(
@@ -78,18 +81,28 @@ def trim_to_layer(
         start=0,
         length=x.size(0) - num_sampled_nodes_per_hop[-layer],
     )
-    edge_index = edge_index.narrow(
-        dim=1,
-        start=0,
-        length=edge_index.size(1) - num_sampled_edges_per_hop[-layer],
-    )
     if edge_attr is not None:
         edge_attr = edge_attr.narrow(
             dim=0,
             start=0,
             length=edge_attr.size(0) - num_sampled_edges_per_hop[-layer],
         )
-    return x, edge_index, edge_attr
+    if isinstance(edge_index, Tensor):
+        edge_index = edge_index.narrow(
+            dim=1,
+            start=0,
+            length=edge_index.size(1) - num_sampled_edges_per_hop[-layer],
+        )
+        return x, edge_index, edge_attr
+
+    elif isinstance(edge_index, SparseTensor):
+        num_nodes = edge_index.size(0) - num_sampled_nodes_per_hop[-layer]
+        num_seed_nodes = num_nodes - num_sampled_nodes_per_hop[-(layer + 1)]
+        edge_index = trim_sparse_tensor(edge_index, num_nodes, num_seed_nodes)
+
+        return x, edge_index, edge_attr
+
+    raise NotImplementedError
 
 
 class TrimToLayer(torch.nn.Module):
@@ -99,7 +112,7 @@ class TrimToLayer(torch.nn.Module):
         num_sampled_nodes_per_hop: Optional[List[int]],
         num_sampled_edges_per_hop: Optional[List[int]],
         x: Tensor,
-        edge_index: Tensor,
+        edge_index: Union[Tensor, SparseTensor],
         edge_attr: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
 
@@ -123,3 +136,51 @@ class TrimToLayer(torch.nn.Module):
             edge_index,
             edge_attr,
         )
+
+
+# Helper functions ############################################################
+
+
+def trim_sparse_tensor(src: SparseTensor, num_nodes: int,
+                       num_seed_nodes: None) -> SparseTensor:
+    r"""Trims a :class:`SparseTensor` along both dimensions to only contain
+    the upper :obj:`num_nodes` in both dimensions.
+
+    It is assumed that :class:`SparseTensor` is obtained from BFS traversing,
+    starting from the nodes that have been initially selected.
+
+    Args:
+        src (SparseTensor): The sparse tensor.
+        num_nodes (int): The number of first nodes to keep.
+        num_seed_nodes (int): The number of seed nodes to compute
+            representations.
+    """
+    rowptr, col, value = src.csr()
+
+    rowptr = torch.narrow(rowptr, 0, 0, num_nodes + 1).clone()
+    rowptr[num_seed_nodes + 1:] = rowptr[num_seed_nodes]
+
+    col = torch.narrow(col, 0, 0, rowptr[-1])
+
+    if value is not None:
+        value = torch.narrow(value, 0, 0, rowptr[-1])
+
+    csr2csc = src.storage._csr2csc
+    if csr2csc is not None:
+        csr2csc = csr2csc[csr2csc < len(col)]
+
+    storage = SparseStorage(
+        row=None,
+        rowptr=rowptr,
+        col=col,
+        value=value,
+        sparse_sizes=(num_nodes, num_nodes),
+        rowcount=None,
+        colptr=None,
+        colcount=None,
+        csr2csc=csr2csc,
+        csc2csr=None,
+        is_sorted=True,
+        trust_data=True,
+    )
+    return src.from_storage(storage)


Test patch:
diff --git a/test/utils/test_trim_to_layer.py b/test/utils/test_trim_to_layer.py
index 64c99165b..3260a133c 100644
--- a/test/utils/test_trim_to_layer.py
+++ b/test/utils/test_trim_to_layer.py
@@ -3,56 +3,126 @@ from typing import List, Optional
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.data import Data
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import GraphConv
 from torch_geometric.testing import withPackage
+from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import trim_to_layer
+from torch_geometric.utils.trim_to_layer import trim_sparse_tensor
+
+
+@withPackage('torch_sparse')
+def test_trim_sparse_tensor():
+    edge_index = torch.tensor([[0, 0, 1, 2], [1, 2, 3, 4]])
+    adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=[5, 5])
+
+    adj = trim_sparse_tensor(adj, num_nodes=3, num_seed_nodes=1)
+
+    row, col, _ = adj.coo()
+    assert row.tolist() == [0, 0]
+    assert col.tolist() == [1, 2]
 
 
 def test_trim_to_layer_basic():
-    x = torch.arange(4)
-    edge_index = torch.tensor([[1, 2, 3], [0, 1, 2]])
-    edge_weight = torch.arange(3)
+    x0 = torch.arange(4)
+    edge_index0 = torch.tensor([[1, 2, 3], [0, 1, 2]])
+    edge_weight0 = torch.arange(3)
 
-    num_sampled_nodes_per_hop = [1, 1, 1, 1]
+    num_sampled_nodes_per_hop = [1, 1, 1]
     num_sampled_edges_per_hop = [1, 1, 1]
 
-    x, edge_index, edge_weight = trim_to_layer(
+    x1, edge_index1, edge_weight1 = trim_to_layer(
         layer=0,
         num_sampled_nodes_per_hop=num_sampled_nodes_per_hop,
         num_sampled_edges_per_hop=num_sampled_edges_per_hop,
-        x=x,
-        edge_index=edge_index,
-        edge_attr=edge_weight,
+        x=x0,
+        edge_index=edge_index0,
+        edge_attr=edge_weight0,
     )
-    assert torch.equal(x, torch.arange(4))
-    assert edge_index.tolist() == [[1, 2, 3], [0, 1, 2]]
-    assert torch.equal(edge_weight, torch.arange(3))
-
-    x, edge_index, edge_weight = trim_to_layer(
+    assert torch.equal(x1, torch.arange(4))
+    assert edge_index1.tolist() == [[1, 2, 3], [0, 1, 2]]
+    assert torch.equal(edge_weight1, torch.arange(3))
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj0 = SparseTensor.from_edge_index(edge_index0, edge_weight0, (4, 4))
+        x1, adj_t1, _ = trim_to_layer(
+            layer=0,
+            num_sampled_nodes_per_hop=num_sampled_nodes_per_hop,
+            num_sampled_edges_per_hop=num_sampled_edges_per_hop,
+            x=x0,
+            edge_index=adj0.t(),
+            edge_attr=edge_weight0,
+        )
+        adj1 = adj_t1.t()
+        assert adj1.sizes() == [4, 4]
+
+        row, col, value = adj1.coo()
+        assert torch.equal(x1, torch.arange(4))
+        assert row.tolist() == [1, 2, 3]
+        assert col.tolist() == [0, 1, 2]
+        assert torch.equal(value, torch.arange(3))
+
+    x2, edge_index2, edge_weight2 = trim_to_layer(
         layer=1,
         num_sampled_nodes_per_hop=num_sampled_nodes_per_hop,
         num_sampled_edges_per_hop=num_sampled_edges_per_hop,
-        x=x,
-        edge_index=edge_index,
-        edge_attr=edge_weight,
+        x=x1,
+        edge_index=edge_index1,
+        edge_attr=edge_weight1,
     )
-    assert torch.equal(x, torch.arange(3))
-    assert edge_index.tolist() == [[1, 2], [0, 1]]
-    assert torch.equal(edge_weight, torch.arange(2))
-
-    x, edge_index, edge_weight = trim_to_layer(
+    assert torch.equal(x2, torch.arange(3))
+    assert edge_index2.tolist() == [[1, 2], [0, 1]]
+    assert torch.equal(edge_weight2, torch.arange(2))
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj1 = SparseTensor.from_edge_index(edge_index1, edge_weight1, (4, 4))
+        x2, adj_t2, _ = trim_to_layer(
+            layer=1,
+            num_sampled_nodes_per_hop=num_sampled_nodes_per_hop,
+            num_sampled_edges_per_hop=num_sampled_edges_per_hop,
+            x=x1,
+            edge_index=adj1.t(),
+        )
+        adj2 = adj_t2.t()
+        assert adj2.sizes() == [3, 3]
+
+        row, col, value = adj2.coo()
+        assert torch.equal(x2, torch.arange(3))
+        assert row.tolist() == [1, 2]
+        assert col.tolist() == [0, 1]
+        assert torch.equal(value, torch.arange(2))
+
+    x3, edge_index3, edge_weight3 = trim_to_layer(
         layer=2,
         num_sampled_nodes_per_hop=num_sampled_nodes_per_hop,
         num_sampled_edges_per_hop=num_sampled_edges_per_hop,
-        x=x,
-        edge_index=edge_index,
-        edge_attr=edge_weight,
+        x=x2,
+        edge_index=edge_index2,
+        edge_attr=edge_weight2,
     )
-    assert torch.equal(x, torch.arange(2))
-    assert edge_index.tolist() == [[1], [0]]
-    assert torch.equal(edge_weight, torch.arange(1))
+    assert torch.equal(x3, torch.arange(2))
+    assert edge_index3.tolist() == [[1], [0]]
+    assert torch.equal(edge_weight3, torch.arange(1))
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index2, edge_weight2, (3, 3))
+        x3, adj_t3, _ = trim_to_layer(
+            layer=2,
+            num_sampled_nodes_per_hop=num_sampled_nodes_per_hop,
+            num_sampled_edges_per_hop=num_sampled_edges_per_hop,
+            x=x2,
+            edge_index=adj2.t(),
+        )
+        adj3 = adj_t3.t()
+        assert adj3.sizes() == [2, 2]
+
+        row, col, value = adj3.coo()
+        assert torch.equal(x3, torch.arange(2))
+        assert row.tolist() == [1]
+        assert col.tolist() == [0]
+        assert torch.equal(value, torch.arange(1))
 
 
 def test_trim_to_layer_hetero():