Fix data.num_edges for torch.sparse.Tensor (#7104)

Fixes zero `num_edges` when using native PyTorch sparse tensor (#7103). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: rusty1s <[email protected]>
pyg-team · Apr 3, 2023 · 3c15018 · 3c15018
1 parent 8555970
commit 3c15018
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Added support for `Data.num_edges` for native `torch.sparse.Tensor` adjacency matrices ([#7104](https://github.com/pyg-team/pytorch_geometric/pull/7104))
 - Fixed crash of heterogeneous data loaders if node or edge types are missing ([#7060](https://github.com/pyg-team/pytorch_geometric/pull/7060), [#7087](https://github.com/pyg-team/pytorch_geometric/pull/7087))
 - Accelerated attention-based `MultiAggregation` ([#7077](https://github.com/pyg-team/pytorch_geometric/pull/7077))
 - Edges in `HeterophilousGraphDataset` are now undirected by default ([#7065](https://github.com/pyg-team/pytorch_geometric/pull/7065))

diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
@@ -26,6 +26,7 @@
 from torch_geometric.utils import (
     coalesce,
     contains_isolated_nodes,
+    is_torch_sparse_tensor,
     is_undirected,
 )
 
@@ -418,6 +419,8 @@ def num_edges(self) -> int:
         for value in self.values('adj', 'adj_t'):
             if isinstance(value, SparseTensor):
                 return value.nnz()
+            elif is_torch_sparse_tensor(value):
+                return value._nnz()
         return 0
 
     @property