Repository: pyg-team/pytorch_geometric
Pull request #7061 (2023-03-28 09:06:03)

Drop `torch_sparse` dependency in tests (10/n)

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

> :exclamation: Current head 407bf7c differs from pull request most recent head 93101ed. Consider uploading reports for the commit 93101ed to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7061   +/-   ##
=======================================
  Coverage   91.43%   91.43%           
=======================================
  Files         436      436           
  Lines       23826    23828    +2     
=======================================
+ Hits        21786    21788    +2     
  Misses       2040     2040           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7061?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/typing.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7061?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3R5cGluZy5weQ==) | `100.00% <ø> (ø)` | |
| [torch\_geometric/nn/conv/message\_passing.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7061?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvbWVzc2FnZV9wYXNzaW5nLnB5) | `97.34% <100.00%> (+<0.01%)` | :arrow_up: |
| [torch\_geometric/nn/conv/utils/typing.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7061?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvdXRpbHMvdHlwaW5nLnB5) | `83.95% <100.00%> (+0.20%)` | :arrow_up: |

: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)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd66199f6..ab6051e61 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 support for `jittable` without `torch_sparse` being installed ([#7061](https://github.com/pyg-team/pytorch_geometric/pull/7061))
 - Added unbatching logic for `torch.sparse` tensors ([#7037](https://github.com/pyg-team/pytorch_geometric/pull/7037))
 - Added the `RotatE` KGE model ([#7026](https://github.com/pyg-team/pytorch_geometric/pull/7026))
 
diff --git a/torch_geometric/nn/conv/message_passing.jinja b/torch_geometric/nn/conv/message_passing.jinja
index 54a8ab0c9..13a32aa4c 100644
--- a/torch_geometric/nn/conv/message_passing.jinja
+++ b/torch_geometric/nn/conv/message_passing.jinja
@@ -4,7 +4,6 @@ from torch_geometric.typing import *
 
 import torch
 from torch import Tensor
-from torch_geometric.typing import SparseTensor, torch_sparse
 from torch_geometric.nn.conv.message_passing import *
 from {{module}} import *
 
@@ -63,11 +62,10 @@ class {{cls_name}}({{parent_cls_name}}):
             return src.index_select(self.node_dim, index)
 
         elif isinstance(edge_index, SparseTensor):
+            row, col, _ = edge_index.coo()
             if dim == 0:
-                col = edge_index.storage.col()
                 return src.index_select(self.node_dim, col)
             elif dim == 1:
-                row = edge_index.storage.row()
                 return src.index_select(self.node_dim, row)
 
         raise ValueError(
@@ -126,20 +124,20 @@ class {{cls_name}}({{parent_cls_name}}):
             edge_index_j = edge_def[j]
         elif isinstance(edge_def, SparseTensor):
             adj_t = edge_def
-            edge_index_i = edge_def.storage.row()
-            edge_index_j = edge_def.storage.col()
-            ptr = edge_def.storage.rowptr()
+            edge_index_i, edge_index_j, value = edge_def.coo()
+            ptr, _, _ = edge_def.csr()
+
             {% if 'edge_weight' in collect_types.keys() %}
             if edge_weight is None:
-                edge_weight = edge_def.storage.value()
+                edge_weight = value
             {% endif %}
             {% if 'edge_attr' in collect_types.keys() %}
             if edge_attr is None:
-                edge_attr = edge_def.storage.value()
+                edge_attr = value
             {% endif %}
             {% if 'edge_type' in collect_types.keys() %}
             if edge_type is None:
-                edge_type = edge_def.storage.value()
+                edge_type = value
             {% endif %}
 
         {% if collect_types.get('edge_weight', 'Optional')[:8] != 'Optional' %}assert edge_weight is not None{% endif %}
@@ -205,20 +203,20 @@ class {{cls_name}}({{parent_cls_name}}):
             edge_index_j = edge_def[j]
         elif isinstance(edge_def, SparseTensor):
             adj_t = edge_def
-            edge_index_i = edge_def.storage.row()
-            edge_index_j = edge_def.storage.col()
-            ptr = edge_def.storage.rowptr()
+            edge_index_i, edge_index_j, value = edge_def.coo()
+            ptr, _, _ = edge_def.csr()
+
             {% if 'edge_weight' in edge_collect_types.keys() %}
             if edge_weight is None:
-                edge_weight = edge_def.storage.value()
+                edge_weight = value
             {% endif %}
             {% if 'edge_attr' in edge_collect_types.keys() %}
             if edge_attr is None:
-                edge_attr = edge_def.storage.value()
+                edge_attr = value
             {% endif %}
             {% if 'edge_type' in edge_collect_types.keys() %}
             if edge_type is None:
-                edge_type = edge_def.storage.value()
+                edge_type = value
             {% endif %}
 
         {% if edge_collect_types.get('edge_weight', 'Optional')[:8] != 'Optional' %}assert edge_weight is not None{% endif %}
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index 2db5c9893..95339bf18 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -1,6 +1,7 @@
 import inspect
 import os
 import os.path as osp
+import random
 import re
 from collections import OrderedDict
 from inspect import Parameter
@@ -15,7 +16,6 @@ from typing import (
     Union,
     get_type_hints,
 )
-from uuid import uuid1
 
 import torch
 from torch import Tensor
@@ -302,11 +302,10 @@ class MessagePassing(torch.nn.Module):
                 raise e
 
         elif isinstance(edge_index, SparseTensor):
+            row, col, _ = edge_index.coo()
             if dim == 0:
-                col = edge_index.storage.col()
                 return src.index_select(self.node_dim, col)
             elif dim == 1:
-                row = edge_index.storage.row()
                 return src.index_select(self.node_dim, row)
 
         raise ValueError(
@@ -359,17 +358,20 @@ class MessagePassing(torch.nn.Module):
             out['ptr'] = None
 
         elif isinstance(edge_index, SparseTensor):
+            row, col, value = edge_index.coo()
+            rowptr, _, _ = edge_index.csr()
+
             out['adj_t'] = edge_index
             out['edge_index'] = None
-            out['edge_index_i'] = edge_index.storage.row()
-            out['edge_index_j'] = edge_index.storage.col()
-            out['ptr'] = edge_index.storage.rowptr()
+            out['edge_index_i'] = row
+            out['edge_index_j'] = col
+            out['ptr'] = rowptr
             if out.get('edge_weight', None) is None:
-                out['edge_weight'] = edge_index.storage.value()
+                out['edge_weight'] = value
             if out.get('edge_attr', None) is None:
-                out['edge_attr'] = edge_index.storage.value()
+                out['edge_attr'] = value
             if out.get('edge_type', None) is None:
-                out['edge_type'] = edge_index.storage.value()
+                out['edge_type'] = value
 
         out['index'] = out['edge_index_i']
         out['size'] = size
@@ -867,7 +869,7 @@ class MessagePassing(torch.nn.Module):
         with open(osp.join(root, 'message_passing.jinja'), 'r') as f:
             template = Template(f.read())
 
-        uid = uuid1().hex[:6]
+        uid = '%06x' % random.randrange(16**6)
         cls_name = f'{self.__class__.__name__}Jittable_{uid}'
         jit_module_repr = template.render(
             uid=uid,
diff --git a/torch_geometric/nn/conv/utils/typing.py b/torch_geometric/nn/conv/utils/typing.py
index 25fdfdd34..f4bcb4003 100644
--- a/torch_geometric/nn/conv/utils/typing.py
+++ b/torch_geometric/nn/conv/utils/typing.py
@@ -27,6 +27,8 @@ def sanitize(type_repr: str):
     type_repr = type_repr.replace('typing.', '')
     type_repr = type_repr.replace('torch_sparse.tensor.', '')
     type_repr = type_repr.replace('Adj', 'Union[Tensor, SparseTensor]')
+    type_repr = type_repr.replace('torch_geometric.SparseTensor',
+                                  'SparseTensor')
 
     # Replace `Union[..., NoneType]` by `Optional[...]`.
     sexp = pp.nestedExpr(opener='[', closer=']')
diff --git a/torch_geometric/nn/sequential.jinja b/torch_geometric/nn/sequential.jinja
index 2e5abc673..013d6f46c 100644
--- a/torch_geometric/nn/sequential.jinja
+++ b/torch_geometric/nn/sequential.jinja
@@ -3,7 +3,6 @@ from torch_geometric.typing import *
 
 import torch
 from torch import Tensor
-from torch_geometric.typing import SparseTensor
 
 
 class {{cls_name}}(torch.nn.Module):
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 75af7e5c2..e7379e7c8 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -73,6 +73,12 @@ except (ImportError, OSError) as e:
         def is_cuda(self) -> bool:
             raise ImportError("'SparseTensor' requires 'torch-sparse'")
 
+        def coo(self) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+            raise ImportError("'SparseTensor' requires 'torch-sparse'")
+
+        def csr(self) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+            raise ImportError("'SparseTensor' requires 'torch-sparse'")
+
         def to_torch_sparse_csr_tensor(
             self,
             dtype: Optional[torch.dtype] = None,


Test patch:
diff --git a/test/nn/conv/test_message_passing.py b/test/nn/conv/test_message_passing.py
index 005192c64..46533b089 100644
--- a/test/nn/conv/test_message_passing.py
+++ b/test/nn/conv/test_message_passing.py
@@ -6,6 +6,7 @@ import torch
 from torch import Tensor
 from torch.nn import Linear
 
+import torch_geometric.typing
 from torch_geometric.nn import MessagePassing, aggr
 from torch_geometric.typing import (
     Adj,
@@ -14,7 +15,7 @@ from torch_geometric.typing import (
     Size,
     SparseTensor,
 )
-from torch_geometric.utils import scatter, spmm
+from torch_geometric.utils import scatter, spmm, to_torch_csc_tensor
 
 
 class MyConv(MessagePassing):
@@ -56,24 +57,28 @@ def test_my_conv_basic():
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    value = torch.randn(row.size(0))
-    adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
-    torch_adj = adj.to_torch_sparse_csc_tensor()
+    value = torch.randn(edge_index.size(1))
+    adj1 = to_torch_csc_tensor(edge_index, value, size=(4, 4))
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index, value, (4, 4))
 
     conv = MyConv(8, 32)
     out = conv(x1, edge_index, value)
     assert out.size() == (4, 32)
     assert torch.allclose(conv(x1, edge_index, value, (4, 4)), out, atol=1e-6)
-    assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
-    assert torch.allclose(conv(x1, torch_adj.t()), out, atol=1e-6)
+    assert torch.allclose(conv(x1, adj1.t()), out, atol=1e-6)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv(x1, adj2.t()), out, atol=1e-6)
     conv.fuse = False
-    assert torch.allclose(conv(x1, adj.t()), out)
-    assert torch.allclose(conv(x1, torch_adj.t()), out, atol=1e-6)
+    assert torch.allclose(conv(x1, adj1.t()), out, atol=1e-6)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv(x1, adj2.t()), out, atol=1e-6)
     conv.fuse = True
 
-    adj = adj.sparse_resize((4, 2))
-    torch_adj = adj.to_torch_sparse_csc_tensor()
+    # Bipartite message passing:
+    adj1 = to_torch_csc_tensor(edge_index, value, size=(4, 2))
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index, value, (4, 2))
 
     conv = MyConv((8, 16), 32)
     out1 = conv((x1, x2), edge_index, value)
@@ -81,19 +86,21 @@ def test_my_conv_basic():
     assert out1.size() == (2, 32)
     assert out2.size() == (2, 32)
     assert torch.allclose(conv((x1, x2), edge_index, value, (4, 2)), out1)
-    assert torch.allclose(conv((x1, x2), adj.t()), out1, atol=1e-6)
-    assert torch.allclose(conv((x1, x2), torch_adj.t()), out1, atol=1e-6)
-    assert torch.allclose(conv((x1, None), adj.t()), out2, atol=1e-6)
-    assert torch.allclose(conv((x1, None), torch_adj.t()), out2, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj1.t()), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, None), adj1.t()), out2, atol=1e-6)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv((x1, x2), adj2.t()), out1, atol=1e-6)
+        assert torch.allclose(conv((x1, None), adj2.t()), out2, atol=1e-6)
     conv.fuse = False
-    assert torch.allclose(conv((x1, x2), adj.t()), out1, atol=1e-6)
-    assert torch.allclose(conv((x1, x2), torch_adj.t()), out1, atol=1e-6)
-    assert torch.allclose(conv((x1, None), adj.t()), out2, atol=1e-6)
-    assert torch.allclose(conv((x1, None), torch_adj.t()), out2, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj1.t()), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, None), adj1.t()), out2, atol=1e-6)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv((x1, x2), adj2.t()), out1, atol=1e-6)
+        assert torch.allclose(conv((x1, None), adj2.t()), out2, atol=1e-6)
 
     # Test gradient computation for `torch.sparse` tensors:
     conv.fuse = True
-    torch_adj_t = torch_adj.t().requires_grad_()
+    torch_adj_t = adj1.t().requires_grad_()
     out = conv((x1, x2), torch_adj_t)
     out.sum().backward()
     assert torch_adj_t.grad is not None
@@ -118,26 +125,29 @@ def test_my_conv_jittable():
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    value = torch.randn(row.size(0))
-    adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
+    value = torch.randn(edge_index.size(1))
 
     conv = MyConv(8, 32)
     out = conv(x1, edge_index, value)
 
+    conv.jittable()  # Should succeed.
+    torch.jit.script(conv.jittable())  # Should succeed.
+
     t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
     jit = torch.jit.script(conv.jittable(t))
     assert torch.allclose(jit(x1, edge_index, value), out, atol=1e-6)
     assert torch.allclose(jit(x1, edge_index, value, (4, 4)), out, atol=1e-6)
 
-    t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
-    jit = torch.jit.script(conv.jittable(t))
-    assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
-    jit.fuse = False
-    assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
-    jit.fuse = True
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, value, (4, 4))
+
+        t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
+        jit = torch.jit.script(conv.jittable(t))
+        assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
+        jit.fuse = False
+        assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
+        jit.fuse = True
 
-    adj = adj.sparse_resize((4, 2))
     conv = MyConv((8, 16), 32)
     out1 = conv((x1, x2), edge_index, value)
     out2 = conv((x1, None), edge_index, value, (4, 2))
@@ -148,14 +158,17 @@ def test_my_conv_jittable():
     assert torch.allclose(jit((x1, x2), edge_index, value, (4, 2)), out1)
     assert torch.allclose(jit((x1, None), edge_index, value, (4, 2)), out2)
 
-    t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
-    jit = torch.jit.script(conv.jittable(t))
-    assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
-    assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
-    jit.fuse = False
-    assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
-    assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
-    jit.fuse = True
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, value, (4, 2))
+
+        t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
+        jit = torch.jit.script(conv.jittable(t))
+        assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
+        assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
+        jit.fuse = False
+        assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
+        assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
+        jit.fuse = True
 
 
 @pytest.mark.parametrize('aggr', ['add', 'sum', 'mean', 'min', 'max', 'mul'])
@@ -173,25 +186,29 @@ def test_my_static_graph_conv():
     x1 = torch.randn(3, 4, 8)
     x2 = torch.randn(3, 2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    value = torch.randn(row.size(0))
-    adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
+    value = torch.randn(edge_index.size(1))
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, value, (4, 4))
 
     conv = MyConv(8, 32)
     out = conv(x1, edge_index, value)
     assert out.size() == (3, 4, 32)
     assert torch.allclose(conv(x1, edge_index, value, (4, 4)), out)
-    assert torch.allclose(conv(x1, adj.t()), out)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv(x1, adj.t()), out)
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, value, (4, 2))
 
-    adj = adj.sparse_resize((4, 2))
     conv = MyConv((8, 16), 32)
     out1 = conv((x1, x2), edge_index, value)
-    assert out1.size() == (3, 2, 32)
-    assert torch.allclose(conv((x1, x2), edge_index, value, (4, 2)), out1)
-    assert torch.allclose(conv((x1, x2), adj.t()), out1)
     out2 = conv((x1, None), edge_index, value, (4, 2))
+    assert out1.size() == (3, 2, 32)
     assert out2.size() == (3, 2, 32)
-    assert torch.allclose(conv((x1, None), adj.t()), out2)
+    assert torch.allclose(conv((x1, x2), edge_index, value, (4, 2)), out1)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv((x1, x2), adj.t()), out1)
+        assert torch.allclose(conv((x1, None), adj.t()), out2)
 
 
 class MyMultipleAggrConv(MessagePassing):
@@ -214,22 +231,21 @@ def test_my_multiple_aggr_conv(multi_aggr_tuple):
     aggr_kwargs, expand = multi_aggr_tuple
     x = torch.randn(4, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
-    torch_adj = adj.to_torch_sparse_csc_tensor()
+    adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
 
     conv = MyMultipleAggrConv(aggr_kwargs=aggr_kwargs)
     out = conv(x, edge_index)
     assert out.size() == (4, 16 * expand)
-    assert torch.allclose(conv(x, adj.t()), out)
-    assert torch.allclose(conv(x, torch_adj.t()), out)
+    assert torch.allclose(conv(x, adj1.t()), out)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        assert torch.allclose(conv(x, adj2.t()), out)
 
 
 def test_my_multiple_aggr_conv_jittable():
     x = torch.randn(4, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
 
     conv = MyMultipleAggrConv()
     out = conv(x, edge_index)
@@ -238,9 +254,11 @@ def test_my_multiple_aggr_conv_jittable():
     jit = torch.jit.script(conv.jittable(t))
     assert torch.allclose(jit(x, edge_index), out)
 
-    t = '(Tensor, SparseTensor) -> Tensor'
-    jit = torch.jit.script(conv.jittable(t))
-    assert torch.allclose(jit(x, adj.t()), out)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        t = '(Tensor, SparseTensor) -> Tensor'
+        jit = torch.jit.script(conv.jittable(t))
+        assert torch.allclose(jit(x, adj.t()), out)
 
 
 def test_copy():
@@ -283,25 +301,25 @@ class MyEdgeConv(MessagePassing):
 def test_my_edge_conv():
     x = torch.randn(4, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
-    torch_adj = adj.to_torch_sparse_csc_tensor()
+    adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
+    row, col = edge_index
     expected = scatter(x[row] - x[col], col, dim=0, dim_size=4, reduce='sum')
 
     conv = MyEdgeConv()
     out = conv(x, edge_index)
     assert out.size() == (4, 16)
     assert torch.allclose(out, expected)
-    assert torch.allclose(conv(x, adj.t()), out)
-    assert torch.allclose(conv(x, torch_adj.t()), out)
+    assert torch.allclose(conv(x, adj1.t()), out)
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        assert torch.allclose(conv(x, adj2.t()), out)
 
 
 def test_my_edge_conv_jittable():
     x = torch.randn(4, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
 
     conv = MyEdgeConv()
     out = conv(x, edge_index)
@@ -310,9 +328,11 @@ def test_my_edge_conv_jittable():
     jit = torch.jit.script(conv.jittable(t))
     assert torch.allclose(jit(x, edge_index), out)
 
-    t = '(Tensor, SparseTensor) -> Tensor'
-    jit = torch.jit.script(conv.jittable(t))
-    assert torch.allclose(jit(x, adj.t()), out)
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        t = '(Tensor, SparseTensor) -> Tensor'
+        jit = torch.jit.script(conv.jittable(t))
+        assert torch.allclose(jit(x, adj.t()), out)
 
 
 num_pre_hook_calls = 0
@@ -324,9 +344,8 @@ def test_message_passing_hooks():
 
     x = torch.randn(4, 8)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    value = torch.randn(row.size(0))
-    adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
+    value = torch.randn(edge_index.size(1))
+    adj = to_torch_csc_tensor(edge_index, value, size=(4, 4))
 
     def pre_hook(module, inputs):
         assert module == conv
@@ -446,20 +465,22 @@ class MyDefaultArgConv(MessagePassing):
 def test_my_default_arg_conv():
     x = torch.randn(4, 1)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    row, col = edge_index
-    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
-    torch_adj = adj.to_torch_sparse_csc_tensor()
+    adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
     conv = MyDefaultArgConv()
     assert conv(x, edge_index).view(-1).tolist() == [0, 0, 0, 0]
-    assert conv(x, adj.t()).view(-1).tolist() == [0, 0, 0, 0]
-    assert conv(x, torch_adj.t()).view(-1).tolist() == [0, 0, 0, 0]
+    assert conv(x, adj1.t()).view(-1).tolist() == [0, 0, 0, 0]
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        assert conv(x, adj2.t()).view(-1).tolist() == [0, 0, 0, 0]
 
 
 def test_my_default_arg_conv_jittable():
     conv = MyDefaultArgConv()
 
-    with pytest.raises(RuntimeError):  # This should not succeed in JIT mode.
+    # This should not succeed in JIT mode.
+    with pytest.raises((RuntimeError, AttributeError)):
         torch.jit.script(conv.jittable())
 
 
@@ -555,13 +576,17 @@ def test_message_passing_with_aggr_module(aggr_module):
     x = torch.randn(4, 8)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
     row, col = edge_index
-    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
+    adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
     conv = MyAggregatorConv(aggr=aggr_module)
     assert isinstance(conv.aggr_module, aggr.Aggregation)
     out = conv(x, edge_index)
     assert out.size(0) == 4 and out.size(1) in {8, 16}
-    assert torch.allclose(conv(x, adj.t()), out)
+    assert torch.allclose(conv(x, adj1.t()), out)
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        assert torch.allclose(conv(x, adj2.t()), out)
 
 
 def test_message_passing_int32_edge_index():