Repository: pyg-team/pytorch_geometric
Pull request #8759 (2024-01-12 23:13:29)

Implement `torch.sparse` and `SparseTensor` support in `propagate.jinja`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8759?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `6 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`5b7111f`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/5b7111f9348d42c7983429da6f6bf84b284e6409?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.41% compared to head [(`79c74dc`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8759?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.40%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8759?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/\_spmm.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8759?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19zcG1tLnB5) | 0.00% | [6 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8759?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    #8759      +/-   ##
==========================================
- Coverage   89.41%   89.40%   -0.01%     
==========================================
  Files         482      482              
  Lines       31213    31209       -4     
==========================================
- Hits        27908    27902       -6     
- Misses       3305     3307       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8759?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/aggr/base.py b/torch_geometric/nn/aggr/base.py
index 2028a2347..f326e7b9d 100644
--- a/torch_geometric/nn/aggr/base.py
+++ b/torch_geometric/nn/aggr/base.py
@@ -3,6 +3,8 @@ from typing import Optional, Tuple
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
+from torch_geometric import is_compiling
 from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.utils import scatter, segment, to_dense_batch
 
@@ -169,11 +171,14 @@ class Aggregation(torch.nn.Module):
                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                dim: int = -2, reduce: str = 'sum') -> Tensor:
 
-        if ptr is not None:
+        if (ptr is not None and torch_geometric.typing.WITH_TORCH_SCATTER
+                and not is_compiling()):
             ptr = expand_left(ptr, dim, dims=x.dim())
             return segment(x, ptr, reduce=reduce)
 
-        assert index is not None
+        if index is None:
+            raise NotImplementedError(
+                "Aggregation requires 'index' to be specified")
         return scatter(x, index, dim, dim_size, reduce)
 
     def to_dense_batch(
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index 6f3dca6dd..6fbbe1e9e 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -189,7 +189,7 @@ class MessagePassing(torch.nn.Module):
             )
 
             self.propagate_jit = module.propagate
-            self._collect_jit = module._collect
+            self.collect_jit = module._collect
 
     def reset_parameters(self) -> None:
         r"""Resets all learnable parameters of the module."""
diff --git a/torch_geometric/nn/conv/propagate.jinja b/torch_geometric/nn/conv/propagate.jinja
index 6f064b686..fa848cea5 100644
--- a/torch_geometric/nn/conv/propagate.jinja
+++ b/torch_geometric/nn/conv/propagate.jinja
@@ -6,7 +6,9 @@ from torch import Tensor
 
 import torch_geometric.typing
 from torch_geometric import is_compiling
-from torch_geometric.utils import is_sparse
+from torch_geometric.utils import is_sparse, is_torch_sparse_tensor
+from torch_geometric.utils.sparse import ptr2index
+from torch_geometric.typing import SparseTensor
 
 from {{module}} import *
 
@@ -27,42 +29,71 @@ def _collect(
 ) -> CollectArgs:
 
     i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
-{% for name in collect_param_dict %}
+
+    # Collect special arguments:
+    if isinstance(edge_index, Tensor) and is_torch_sparse_tensor(edge_index):
+{%- if 'edge_index' in collect_param_dict %}
+        raise ValueError("Cannot collect 'edge_indices' for sparse matrices")
+{%- endif %}
+        adj_t = edge_index
+        if adj_t.layout == torch.sparse_coo:
+            edge_index_i = adj_t.indices()[0]
+            edge_index_j = adj_t.indices()[1]
+            ptr = None
+        elif adj_t.layout == torch.sparse_csr:
+            ptr = adj_t.crow_indices()
+            edge_index_j = adj_t.col_indices()
+            edge_index_i = ptr2index(ptr, output_size=edge_index_j.numel())
+        else:
+            raise ValueError(f"Recieved invalid layout '{adj_t.layout}'")
+
+    elif isinstance(edge_index, SparseTensor):
+{%- if 'edge_index' in collect_param_dict %}
+        raise ValueError("Cannot collect 'edge_indices' for sparse matrices")
+{%- endif %}
+        adj_t = edge_index
+        edge_index_i, edge_index_j, _value = adj_t.coo()
+        ptr, _, _ = adj_t.csr()
+
+    elif isinstance(edge_index, Tensor):
+{%- if 'adj_t' in collect_param_dict %}
+        raise ValueError("Cannot collect 'adj_t' for edge indices")
+{%- endif %}
+        edge_index_i = edge_index[i]
+        edge_index_j = edge_index[j]
+        ptr = None
+
+    # Collect user-defined arguments:
+{%- for name in collect_param_dict %}
 {%- if name.endswith('_i') or name.endswith('_j') %}
-    # Collect `{{name}}`:
+    # ({{loop.index}}) - Collect `{{name}}`:
     if isinstance({{name[:-2]}}, (tuple, list)):
         assert len({{name[:-2]}}) == 2
         _{{name[:-2]}}_0, _{{name[:-2]}}_1 = {{name[:-2]}}[0], {{name[:-2]}}[1]
         if isinstance(_{{name[:-2]}}_0, Tensor):
             self._set_size(size, 0, _{{name[:-2]}}_0)
 {%- if name.endswith('_j') %}
-            {{name}} = self._lift(_{{name[:-2]}}_0, edge_index, {{name[-1]}})
+            {{name}} = _{{name[:-2]}}_0.index_select(self.node_dim, edge_index_{{name[-1]}})
         else:
             {{name}} = None
 {%- endif %}
         if isinstance(_{{name[:-2]}}_1, Tensor):
             self._set_size(size, 1, _{{name[:-2]}}_1)
 {%- if name.endswith('_i') %}
-            {{name}} = self._lift(_{{name[:-2]}}_1, edge_index, {{name[-1]}})
+            {{name}} = _{{name[:-2]}}_1.index_select(self.node_dim, edge_index_{{name[-1]}})
         else:
             {{name}} = None
 {%- endif %}
     elif isinstance({{name[:-2]}}, Tensor):
         self._set_size(size, 0, {{name[:-2]}})
         self._set_size(size, 1, {{name[:-2]}})
-        {{name}} = self._lift({{name[:-2]}}, edge_index, {{name[-1]}})
+        {{name}} = {{name[:-2]}}.index_select(self.node_dim, edge_index_{{name[-1]}})
     else:
         {{name}} = None
 {%- endif %}
 {%- endfor %}
 
-    assert isinstance(edge_index, Tensor)
-    adj_t = None
-    edge_index_i = edge_index[i]
-    edge_index_j = edge_index[j]
-    ptr = None
     index = edge_index_i
-
     size_i = size[i]
     size_j = size[j]
     dim_size = size_i
@@ -105,7 +136,6 @@ def propagate(
     fuse = is_sparse(edge_index) and self.fuse and self.explain is not True
 
     if fuse:
-        assert isinstance(edge_index, SparseTensor)  # TODO
 
         # Begin Message and Aggregate Forward Pre Hook #########################
         if not torch.jit.is_scripting() and not is_compiling():
@@ -150,6 +180,7 @@ def propagate(
         )
 
     else:
+
         kwargs = self._collect(
             edge_index,
 {%- for name in propagate_signature.param_dict %}
diff --git a/torch_geometric/nn/conv/sage_conv.py b/torch_geometric/nn/conv/sage_conv.py
index e47612be7..bab94159a 100644
--- a/torch_geometric/nn/conv/sage_conv.py
+++ b/torch_geometric/nn/conv/sage_conv.py
@@ -146,8 +146,7 @@ class SAGEConv(MessagePassing):
     def message(self, x_j: Tensor) -> Tensor:
         return x_j
 
-    def message_and_aggregate(self, adj_t: SparseTensor,
-                              x: OptPairTensor) -> Tensor:
+    def message_and_aggregate(self, adj_t: Adj, x: OptPairTensor) -> Tensor:
         if isinstance(adj_t, SparseTensor):
             adj_t = adj_t.set_value(None, layout=None)
         return spmm(adj_t, x[0], reduce=self.aggr)
diff --git a/torch_geometric/utils/_softmax.py b/torch_geometric/utils/_softmax.py
index 19cde3b08..65b43ec48 100644
--- a/torch_geometric/utils/_softmax.py
+++ b/torch_geometric/utils/_softmax.py
@@ -52,10 +52,13 @@ def softmax(
                 [0.7607, 0.2393, 1.0000, 1.0000],
                 [0.8062, 0.1938, 1.0000, 1.0000]])
     """
-    if ptr is not None:
-        if (src.device.type == 'cpu' and torch_geometric.typing.WITH_SOFTMAX
-                and not is_compiling()):  # pragma: no cover
-            return pyg_lib.ops.softmax_csr(src, ptr, dim)
+    if (ptr is not None and src.device.type == 'cpu'
+            and torch_geometric.typing.WITH_SOFTMAX
+            and not is_compiling()):  # pragma: no cover
+        return pyg_lib.ops.softmax_csr(src, ptr, dim)
+
+    if (ptr is not None and torch_geometric.typing.WITH_TORCH_SCATTER
+            and not is_compiling()):
         dim = dim + src.dim() if dim < 0 else dim
         size = ([1] * dim) + [-1]
         count = ptr[1:] - ptr[:-1]
@@ -73,6 +76,6 @@ def softmax(
         out_sum = scatter(out, index, dim, dim_size=N, reduce='sum') + 1e-16
         out_sum = out_sum.index_select(dim, index)
     else:
-        raise NotImplementedError
+        raise NotImplementedError("'softmax' requires 'index' to be specified")
 
     return out / out_sum
diff --git a/torch_geometric/utils/_spmm.py b/torch_geometric/utils/_spmm.py
index c48d01e8b..43d046bf9 100644
--- a/torch_geometric/utils/_spmm.py
+++ b/torch_geometric/utils/_spmm.py
@@ -1,4 +1,3 @@
-import typing
 import warnings
 
 import torch
@@ -8,34 +7,11 @@ import torch_geometric.typing
 from torch_geometric.typing import Adj, SparseTensor, torch_sparse
 from torch_geometric.utils import is_torch_sparse_tensor, scatter
 
-if typing.TYPE_CHECKING:
-    from typing import overload
-else:
-    from torch.jit import _overload as overload
 
-
-@overload
 def spmm(
-    src: Tensor,
+    src: Adj,
     other: Tensor,
     reduce: str = 'sum',
-) -> Tensor:
-    pass
-
-
-@overload
-def spmm(  # noqa: F811
-    src: SparseTensor,
-    other: Tensor,
-    reduce: str = 'sum',
-) -> Tensor:
-    pass
-
-
-def spmm(  # noqa: F811
-        src: Adj,
-        other: Tensor,
-        reduce: str = 'sum',
 ) -> Tensor:
     r"""Matrix product of sparse matrix with dense matrix.
 
@@ -136,14 +112,15 @@ def spmm(  # noqa: F811
         elif (torch_geometric.typing.WITH_PT112
               and src.layout == torch.sparse_csc):
             assert src.layout == torch.sparse_csc
-            deg = scatter(torch.ones_like(src.values()), src.row_indices(),
-                          dim=0, dim_size=src.size(0), reduce='sum')
+            ones = torch.ones_like(src.values())
+            index = src.row_indices()
+            deg = scatter(ones, index, 0, dim_size=src.size(0), reduce='sum')
         else:
             assert src.layout == torch.sparse_coo
             src = src.coalesce()
-            deg = scatter(torch.ones_like(src.values()),
-                          src.indices()[0], dim=0, dim_size=src.size(0),
-                          reduce='sum')
+            ones = torch.ones_like(src.values())
+            index = src.indices()[0]
+            deg = scatter(ones, index, 0, dim_size=src.size(0), reduce='sum')
 
         return torch.sparse.mm(src, other) / deg.view(-1, 1).clamp_(min=1)
 
diff --git a/torch_geometric/utils/num_nodes.py b/torch_geometric/utils/num_nodes.py
index 7142a045c..1f5e3b6b3 100644
--- a/torch_geometric/utils/num_nodes.py
+++ b/torch_geometric/utils/num_nodes.py
@@ -1,4 +1,3 @@
-import typing
 from copy import copy
 from typing import Dict, Optional, Tuple, Union
 
@@ -9,37 +8,8 @@ import torch_geometric
 from torch_geometric import EdgeIndex
 from torch_geometric.typing import EdgeType, NodeType, SparseTensor
 
-if typing.TYPE_CHECKING:
-    from typing import overload
-else:
-    from torch.jit import _overload as overload
 
-
-@overload
 def maybe_num_nodes(
-    edge_index: Tensor,
-    num_nodes: Optional[int] = None,
-) -> int:
-    pass
-
-
-@overload
-def maybe_num_nodes(  # noqa: F811
-    edge_index: Tuple[Tensor, Tensor],
-    num_nodes: Optional[int] = None,
-) -> int:
-    pass
-
-
-@overload
-def maybe_num_nodes(  # noqa: F811
-    edge_index: SparseTensor,
-    num_nodes: Optional[int] = None,
-) -> int:
-    pass
-
-
-def maybe_num_nodes(  # noqa: F811
     edge_index: Union[Tensor, Tuple[Tensor, Tensor], SparseTensor],
     num_nodes: Optional[int] = None,
 ) -> int:
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index 1dabc0cca..0adae68a0 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -483,9 +483,9 @@ def set_sparse_value(adj: Tensor, value: Tensor) -> Tensor:
     raise ValueError(f"Unexpected sparse tensor layout (got '{adj.layout}')")
 
 
-def ptr2index(ptr: Tensor) -> Tensor:
-    ind = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
-    return ind.repeat_interleave(ptr[1:] - ptr[:-1])
+def ptr2index(ptr: Tensor, output_size: Optional[int] = None) -> Tensor:
+    index = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
+    return index.repeat_interleave(ptr.diff(), output_size=output_size)
 
 
 def index2ptr(index: Tensor, size: Optional[int] = None) -> Tensor:


Test patch:
diff --git a/test/nn/aggr/test_attention.py b/test/nn/aggr/test_attention.py
index 17ce041d0..f3e9e9d9c 100644
--- a/test/nn/aggr/test_attention.py
+++ b/test/nn/aggr/test_attention.py
@@ -23,7 +23,7 @@ def test_attentional_aggregation():
     assert out.size() == (3, channels)
 
     if not torch_geometric.typing.WITH_TORCH_SCATTER:
-        with pytest.raises(ImportError, match="'segment' requires"):
+        with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
diff --git a/test/nn/aggr/test_basic.py b/test/nn/aggr/test_basic.py
index 07841ec20..8d931b52f 100644
--- a/test/nn/aggr/test_basic.py
+++ b/test/nn/aggr/test_basic.py
@@ -56,7 +56,7 @@ def test_basic_aggregation(Aggregation):
         with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
     elif not torch_geometric.typing.WITH_TORCH_SCATTER:
-        with pytest.raises(ImportError, match="'segment' requires"):
+        with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
@@ -102,7 +102,7 @@ def test_learnable_aggregation(Aggregation, learn):
     assert out.size() == (3, x.size(1))
 
     if not torch_geometric.typing.WITH_TORCH_SCATTER:
-        with pytest.raises(ImportError, match="'segment' requires"):
+        with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
@@ -129,7 +129,7 @@ def test_learnable_channels_aggregation(Aggregation):
     assert out.size() == (3, x.size(1))
 
     if not torch_geometric.typing.WITH_TORCH_SCATTER:
-        with pytest.raises(ImportError, match="'segment' requires"):
+        with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
diff --git a/test/nn/aggr/test_multi.py b/test/nn/aggr/test_multi.py
index 222696b62..3b88f1568 100644
--- a/test/nn/aggr/test_multi.py
+++ b/test/nn/aggr/test_multi.py
@@ -40,7 +40,7 @@ def test_multi_aggr(multi_aggr_tuple):
     assert out.size() == (4, expand * x.size(1))
 
     if not torch_geometric.typing.WITH_TORCH_SCATTER:
-        with pytest.raises(ImportError, match="'segment' requires"):
+        with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
diff --git a/test/nn/aggr/test_scaler.py b/test/nn/aggr/test_scaler.py
index 7cbf6db65..ce88a457c 100644
--- a/test/nn/aggr/test_scaler.py
+++ b/test/nn/aggr/test_scaler.py
@@ -22,5 +22,5 @@ def test_degree_scaler_aggregation(train_norm):
     assert out.size() == (3, 240)
     assert torch.allclose(torch.jit.script(aggr)(x, index), out)
 
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(NotImplementedError, match="requires 'index'"):
         aggr(x, ptr=ptr)
diff --git a/test/nn/conv/test_propagate.py b/test/nn/conv/test_propagate.py
index 3f3f12f96..b5b7525cd 100644
--- a/test/nn/conv/test_propagate.py
+++ b/test/nn/conv/test_propagate.py
@@ -1,11 +1,15 @@
+import pytest
 import torch
 
 from torch_geometric.nn import SAGEConv
+from torch_geometric.utils import to_torch_csc_tensor
 
 
-def test_propagate():
+@pytest.mark.parametrize('fuse', [True, False])
+def test_propagate(fuse):
     SAGEConv.jit_on_init = True
-    conv = SAGEConv(10, 16)
+    conv = SAGEConv(16, 32)
+    conv.fuse = fuse
     SAGEConv.jit_on_init = False
 
     old_propagate = conv.__class__.propagate
@@ -13,8 +17,17 @@ def test_propagate():
 
     try:
         conv.__class__.propagate = conv.propagate_jit
-        conv.__class__._collect = conv._collect_jit
-        torch.jit.script(conv)
+        conv.__class__._collect = conv.collect_jit
+        jit = torch.jit.script(conv)
+
+        x = torch.randn(4, 16)
+        edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
+        adj = to_torch_csc_tensor(edge_index, size=(4, 4))
+
+        out1 = jit(x, edge_index)
+        out2 = jit(x, adj.t())
+        assert torch.allclose(out1, out2)
+
     finally:
         conv.__class__.propagate = old_propagate
         conv.__class__._collect = old_collect
diff --git a/test/utils/test_softmax.py b/test/utils/test_softmax.py
index 731a389bc..817144591 100644
--- a/test/utils/test_softmax.py
+++ b/test/utils/test_softmax.py
@@ -17,10 +17,10 @@ def test_softmax():
     out = softmax(src, index)
     assert out.tolist() == [0.5, 0.5, 1, 1]
     if CALCULATION_VIA_PTR_AVAILABLE:
-        assert softmax(src, None, ptr).tolist() == out.tolist()
+        assert softmax(src, ptr=ptr).tolist() == out.tolist()
     else:
-        with pytest.raises(ImportError):
-            softmax(src, None, ptr)
+        with pytest.raises(NotImplementedError, match="requires 'index'"):
+            softmax(src, ptr=ptr)
 
     src = src.view(-1, 1)
     out = softmax(src, index)