Repository: pyg-team/pytorch_geometric
Pull request #8767 (2024-01-15 12:32:46)

Introduce `edge_updater.jinja`


Patch:
diff --git a/torch_geometric/nn/conv/collect.jinja b/torch_geometric/nn/conv/collect.jinja
new file mode 100644
index 000000000..0ccced855
--- /dev/null
+++ b/torch_geometric/nn/conv/collect.jinja
@@ -0,0 +1,122 @@
+from typing import List, NamedTuple, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import is_torch_sparse_tensor
+from torch_geometric.utils.sparse import ptr2index
+from torch_geometric.typing import SparseTensor
+
+
+class CollectArgs(NamedTuple):
+{%- if collect_param_dict|length > 0 %}
+{%- for param in collect_param_dict.values() %}
+    {{param.name}}: {{param.type_repr}}
+{%- endfor %}
+{%- else %}
+    pass
+{%- endif %}
+
+
+def {{collect_name}}(
+    self,
+    edge_index: Union[Tensor, SparseTensor],
+{%- for param in signature.param_dict.values() %}
+    {{param.name}}: {{param.type_repr}},
+{%- endfor %}
+    size: List[Optional[int]],
+) -> CollectArgs:
+
+    i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
+
+    # 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"Received invalid layout '{adj_t.layout}'")
+
+{%- if 'edge_weight' in collect_param_dict %}
+        edge_weight = adj_t.values()
+{%- elif 'edge_attr' in collect_param_dict %}
+        _value = adj_t.values()
+        edge_attr = None if _value.dim() == 1 else _value
+{%- elif 'edge_type' in collect_param_dict %}
+        edge_type = adj_t.values()
+{%- endif %}
+
+    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()
+
+{%- if 'edge_weight' in collect_param_dict %}
+        edge_weight = _value
+{%- elif 'edge_attr' in collect_param_dict %}
+        edge_attr = None if _value is None or _value.dim() == 1 else _value
+{%- elif 'edge_type' in collect_param_dict %}
+        edge_type = _value
+{%- endif %}
+
+    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')) and
+        name not in ['edge_index_i', 'edge_index_j', 'size_i', 'size_j'] %}
+    # ({{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._index_select(_{{name[:-2]}}_0, 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._index_select(_{{name[:-2]}}_1, 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._index_select({{name[:-2]}}, edge_index_{{name[-1]}})
+    else:
+        {{name}} = None
+{%- endif %}
+{%- endfor %}
+
+    index = edge_index_i
+    size_i = size[i]
+    size_j = size[j]
+    dim_size = size_i
+
+    return CollectArgs(
+{%- for name in collect_param_dict %}
+        {{name}},
+{%- endfor %}
+    )
diff --git a/torch_geometric/nn/conv/edge_updater.jinja b/torch_geometric/nn/conv/edge_updater.jinja
new file mode 100644
index 000000000..b9da10ad5
--- /dev/null
+++ b/torch_geometric/nn/conv/edge_updater.jinja
@@ -0,0 +1,71 @@
+import typing
+from typing import Union
+
+import torch
+from torch import Tensor
+
+import torch_geometric.typing
+from torch_geometric import is_compiling
+from torch_geometric.utils import is_sparse
+from torch_geometric.typing import Size, SparseTensor
+
+from {{module}} import *
+
+
+{% include "collect.jinja" %}
+
+
+def edge_updater(
+    self,
+    edge_index: Union[Tensor, SparseTensor],
+{%- for param in signature.param_dict.values() %}
+    {{param.name}}: {{param.type_repr}},
+{%- endfor %}
+    size: Size = None,
+) -> {{signature.return_type_repr}}:
+
+    mutable_size = self._check_input(edge_index, size)
+
+    kwargs = self.{{collect_name}}(
+        edge_index,
+{%- for name in signature.param_dict %}
+        {{name}},
+{%- endfor %}
+        mutable_size,
+    )
+
+    # Begin Edge Update Forward Pre Hook #######################################
+    if not torch.jit.is_scripting() and not is_compiling():
+        for hook in self._edge_update_forward_pre_hooks.values():
+            hook_kwargs = dict(
+{%- for name in collect_param_dict %}
+                {{name}}={{name}},
+{%- endfor %}
+            )
+            res = hook(self, (edge_index, size, hook_kwargs))
+            if res is not None:
+                edge_index, size, hook_kwargs = res
+{%- for name in signature.param_dict %}
+                kwargs.{{name}} = hook_kwargs['{{name}}']
+{%- endfor %}
+    # End Edge Update Forward Pre Hook #########################################
+
+    out = self.edge_update(
+{%- for name in collect_param_dict %}
+        {{name}}=kwargs.{{name}},
+{%- endfor %}
+    )
+
+    # Begin Edge Update Forward Hook ###########################################
+    if not torch.jit.is_scripting() and not is_compiling():
+        for hook in self._edge_update_forward_hooks.values():
+            hook_kwargs = dict(
+{%- for name in collect_param_dict %}
+                {{name}}=kwargs.{{name}},
+{%- endfor %}
+            )
+            res = hook(self, (edge_index, size, hook_kwargs), out)
+            out = res if res is not None else out
+    # End Edge Update Forward Hook #############################################
+
+    return out
diff --git a/torch_geometric/nn/conv/gat_conv.py b/torch_geometric/nn/conv/gat_conv.py
index 53d68869b..dff818ef2 100644
--- a/torch_geometric/nn/conv/gat_conv.py
+++ b/torch_geometric/nn/conv/gat_conv.py
@@ -334,7 +334,8 @@ class GATConv(MessagePassing):
                         "'edge_index' in a 'SparseTensor' form")
 
         # edge_updater_type: (alpha: OptPairTensor, edge_attr: OptTensor)
-        alpha = self.edge_updater(edge_index, alpha=alpha, edge_attr=edge_attr)
+        alpha = self.edge_updater(edge_index, alpha=alpha, edge_attr=edge_attr,
+                                  size=size)
 
         # propagate_type: (x: OptPairTensor, alpha: Tensor)
         out = self.propagate(edge_index, x=x, alpha=alpha, size=size)
@@ -362,7 +363,7 @@ class GATConv(MessagePassing):
 
     def edge_update(self, alpha_j: Tensor, alpha_i: OptTensor,
                     edge_attr: OptTensor, index: Tensor, ptr: OptTensor,
-                    size_i: Optional[int]) -> Tensor:
+                    dim_size: Optional[int]) -> Tensor:
         # Given edge-level attention coefficients for source and target nodes,
         # we simply need to sum them up to "emulate" concatenation:
         alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
@@ -377,7 +378,7 @@ class GATConv(MessagePassing):
             alpha = alpha + alpha_edge
 
         alpha = F.leaky_relu(alpha, self.negative_slope)
-        alpha = softmax(alpha, index, ptr, size_i)
+        alpha = softmax(alpha, index, ptr, dim_size)
         alpha = F.dropout(alpha, p=self.dropout, training=self.training)
         return alpha
 
diff --git a/torch_geometric/nn/conv/gatv2_conv.py b/torch_geometric/nn/conv/gatv2_conv.py
index 1a2a3a8df..21b7771a5 100644
--- a/torch_geometric/nn/conv/gatv2_conv.py
+++ b/torch_geometric/nn/conv/gatv2_conv.py
@@ -324,7 +324,7 @@ class GATv2Conv(MessagePassing):
 
     def edge_update(self, x_j: Tensor, x_i: Tensor, edge_attr: OptTensor,
                     index: Tensor, ptr: OptTensor,
-                    size_i: Optional[int]) -> Tensor:
+                    dim_size: Optional[int]) -> Tensor:
         x = x_i + x_j
 
         if edge_attr is not None:
@@ -337,7 +337,7 @@ class GATv2Conv(MessagePassing):
 
         x = F.leaky_relu(x, self.negative_slope)
         alpha = (x * self.att).sum(dim=-1)
-        alpha = softmax(alpha, index, ptr, size_i)
+        alpha = softmax(alpha, index, ptr, dim_size)
         alpha = F.dropout(alpha, p=self.dropout, training=self.training)
         return alpha
 
diff --git a/torch_geometric/nn/conv/message_passing.jinja b/torch_geometric/nn/conv/message_passing.jinja
index 5beb930af..e5999dcae 100644
--- a/torch_geometric/nn/conv/message_passing.jinja
+++ b/torch_geometric/nn/conv/message_passing.jinja
@@ -216,9 +216,10 @@ class {{cls_name}}({{parent_cls_name}}):
         {% for name, type_hint in edge_updater_types.items() %}
         {{ name }}: {{ type_hint }},
         {% endfor %}
+        size: Size=None,
     ) -> {{ edge_updater_return_type }}:
 
-        the_size = self._check_input(edge_index, size=None)
+        the_size = self._check_input(edge_index, size)
         in_kwargs = EdgeUpdater_{{uid}}({% for k in edge_updater_types.keys() %}{{k}}={{k}}{{ ", " if not loop.last }}{% endfor %})
         kwargs = self._collect_edge(edge_index, the_size, in_kwargs)
         return self.edge_update({% for k in edge_update_args %}{{k}}=kwargs.{{k}}{{ ", " if not loop.last }}{% endfor %})
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index d9b649f8b..95739c577 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -175,11 +175,14 @@ class MessagePassing(torch.nn.Module):
         # Test code for performing on-the-fly TorchScript support:
         if getattr(self, 'jit_on_init', False) and self.decomposed_layers == 1:
             root_dir = osp.dirname(osp.realpath(__file__))
+            base_module = f'{self.__module__}_{self.__class__.__name__}'
             module = module_from_template(
-                module_name=f'{self.__module__}_propagate',
-                module=self.__module__,
+                module_name=f'{base_module}_propagate',
                 template_path=osp.join(root_dir, 'propagate.jinja'),
-                propagate_signature=self._get_propagate_signature(),
+                # Keyword arguments:
+                module=self.__module__,
+                collect_name='collect',
+                signature=self._get_propagate_signature(),
                 collect_param_dict=self.inspector.get_flat_param_dict(
                     ['message', 'aggregate', 'update']),
                 message_args=self.inspector.get_param_names('message'),
@@ -188,9 +191,21 @@ class MessagePassing(torch.nn.Module):
                     'message_and_aggregate'),
                 update_args=self.inspector.get_param_names('update'),
             )
-
             self.propagate_jit = module.propagate
-            self.collect_jit = module._collect
+            self.collect_jit = module.collect
+
+            module = module_from_template(
+                module_name=f'{base_module}_edge_updater',
+                template_path=osp.join(root_dir, 'edge_updater.jinja'),
+                # Keyword arguments:
+                module=self.__module__,
+                collect_name='edge_collect',
+                signature=self._get_edge_updater_signature(),
+                collect_param_dict=self.inspector.get_param_dict(
+                    'edge_update'),
+            )
+            self.edge_updater_jit = module.edge_updater
+            self.edge_collect_jit = module.edge_collect
 
     def reset_parameters(self) -> None:
         r"""Resets all learnable parameters of the module."""
@@ -607,7 +622,12 @@ class MessagePassing(torch.nn.Module):
 
     # Edge-level Updates ######################################################
 
-    def edge_updater(self, edge_index: Adj, **kwargs: Any) -> Tensor:
+    def edge_updater(
+        self,
+        edge_index: Adj,
+        size: Size = None,
+        **kwargs: Any,
+    ) -> Tensor:
         r"""The initial call to compute or update features for each edge in the
         graph.
 
@@ -617,13 +637,21 @@ class MessagePassing(torch.nn.Module):
                 :class:`torch.sparse.Tensor` that defines the underlying graph
                 connectivity/message passing flow.
                 See :meth:`propagate` for more information.
+            size ((int, int), optional): The size :obj:`(N, M)` of the
+                assignment matrix in case :obj:`edge_index` is a
+                :class:`torch.Tensor`.
+                If set to :obj:`None`, the size will be automatically inferred
+                and assumed to be quadratic.
+                This argument is ignored in case :obj:`edge_index` is a
+                :class:`torch_sparse.SparseTensor` or
+                a :class:`torch.sparse.Tensor`. (default: :obj:`None`)
             **kwargs: Any additional data which is needed to compute or update
                 features for each edge in the graph.
         """
         for hook in self._edge_update_forward_pre_hooks.values():
-            res = hook(self, (edge_index, kwargs))
+            res = hook(self, (edge_index, size, kwargs))
             if res is not None:
-                edge_index, kwargs = res
+                edge_index, size, kwargs = res
 
         mutable_size = self._check_input(edge_index, size=None)
 
@@ -635,7 +663,7 @@ class MessagePassing(torch.nn.Module):
         out = self.edge_update(**edge_kwargs)
 
         for hook in self._edge_update_forward_hooks.values():
-            res = hook(self, (edge_index, kwargs), out)
+            res = hook(self, (edge_index, size, kwargs), out)
             if res is not None:
                 out = res
 
diff --git a/torch_geometric/nn/conv/propagate.jinja b/torch_geometric/nn/conv/propagate.jinja
index 07b9c771e..b5bac5429 100644
--- a/torch_geometric/nn/conv/propagate.jinja
+++ b/torch_geometric/nn/conv/propagate.jinja
@@ -1,134 +1,41 @@
 import typing
-from typing import NamedTuple
+from typing import Union
 
 import torch
 from torch import Tensor
 
 import torch_geometric.typing
 from torch_geometric import is_compiling
-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 torch_geometric.utils import is_sparse
+from torch_geometric.typing import Size, SparseTensor
 
 from {{module}} import *
 
 
-class CollectArgs(NamedTuple):
-{%- for param in collect_param_dict.values() %}
-    {{param.name}}: {{param.type_repr}}
-{%- endfor %}
-
-
-def _collect(
-    self,
-    edge_index: Union[Tensor, SparseTensor],
-{%- for param in propagate_signature.param_dict.values() %}
-    {{param.name}}: {{param.type_repr}},
-{%- endfor %}
-    size: List[Optional[int]],
-) -> CollectArgs:
-
-    i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
-
-    # 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')) and
-        name not in ['edge_index_i', 'edge_index_j', 'size_i', 'size_j'] %}
-    # ({{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._index_select(_{{name[:-2]}}_0, 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._index_select(_{{name[:-2]}}_1, 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._index_select(_{{name[:-2]}}, edge_index_{{name[-1]}})
-    else:
-        {{name}} = None
-{%- endif %}
-{%- endfor %}
-
-    index = edge_index_i
-    size_i = size[i]
-    size_j = size[j]
-    dim_size = size_i
-
-    return CollectArgs(
-{%- for name in collect_param_dict %}
-        {{name}},
-{%- endfor %}
-    )
+{% include "collect.jinja" %}
 
 
 def propagate(
     self,
     edge_index: Union[Tensor, SparseTensor],
-{%- for param in propagate_signature.param_dict.values() %}
+{%- for param in signature.param_dict.values() %}
     {{param.name}}: {{param.type_repr}},
 {%- endfor %}
     size: Size = None,
-) -> {{propagate_signature.return_type_repr}}:
-
-    assert self.decomposed_layers == 1  # TODO
+) -> {{signature.return_type_repr}}:
 
     # Begin Propagate Forward Pre Hook #########################################
     if not torch.jit.is_scripting() and not is_compiling():
         for hook in self._propagate_forward_pre_hooks.values():
             hook_kwargs = dict(
-{%- for name in propagate_signature.param_dict %}
+{%- for name in signature.param_dict %}
                 {{name}}={{name}},
 {%- endfor %}
             )
             res = hook(self, (edge_index, size, hook_kwargs))
             if res is not None:
                 edge_index, size, hook_kwargs = res
-{%- for name in propagate_signature.param_dict %}
+{%- for name in signature.param_dict %}
                 {{name}} = hook_kwargs['{{name}}']
 {%- endfor %}
     # End Propagate Forward Pre Hook ###########################################
@@ -186,9 +93,9 @@ def propagate(
 
     else:
 
-        kwargs = self._collect(
+        kwargs = self.{{collect_name}}(
             edge_index,
-{%- for name in propagate_signature.param_dict %}
+{%- for name in signature.param_dict %}
             {{name}},
 {%- endfor %}
             mutable_size,
diff --git a/torch_geometric/nn/conv/propagate.py b/torch_geometric/nn/conv/propagate.py
index 98a84e3be..653b61d5b 100644
--- a/torch_geometric/nn/conv/propagate.py
+++ b/torch_geometric/nn/conv/propagate.py
@@ -4,7 +4,7 @@ import os.path as osp
 import sys
 from typing import Any
 
-from jinja2 import Template
+from jinja2 import Environment, FileSystemLoader
 
 from torch_geometric import get_home_dir
 
@@ -18,8 +18,8 @@ def module_from_template(
     if module_name in sys.modules:  # If module is already loaded, return it:
         return sys.modules[module_name]
 
-    with open(template_path, 'r') as f:
-        template = Template(f.read())
+    env = Environment(loader=FileSystemLoader(osp.dirname(template_path)))
+    template = env.get_template(osp.basename(template_path))
     module_repr = template.render(**kwargs)
 
     instance_dir = osp.join(get_home_dir(), 'propagate')


Test patch:
diff --git a/test/nn/conv/test_propagate.py b/test/nn/conv/test_propagate.py
index b5b7525cd..dc80659a3 100644
--- a/test/nn/conv/test_propagate.py
+++ b/test/nn/conv/test_propagate.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from torch_geometric.nn import SAGEConv
+from torch_geometric.nn import GATConv, SAGEConv
 from torch_geometric.utils import to_torch_csc_tensor
 
 
@@ -10,14 +10,13 @@ def test_propagate(fuse):
     SAGEConv.jit_on_init = True
     conv = SAGEConv(16, 32)
     conv.fuse = fuse
-    SAGEConv.jit_on_init = False
+    conv.__class__.jit_on_init = False
 
     old_propagate = conv.__class__.propagate
-    old_collect = conv.__class__._collect
 
     try:
         conv.__class__.propagate = conv.propagate_jit
-        conv.__class__._collect = conv.collect_jit
+        conv.__class__.collect = conv.collect_jit
         jit = torch.jit.script(conv)
 
         x = torch.randn(4, 16)
@@ -30,4 +29,31 @@ def test_propagate(fuse):
 
     finally:
         conv.__class__.propagate = old_propagate
-        conv.__class__._collect = old_collect
+
+
+def test_edge_updater():
+    GATConv.jit_on_init = True
+    conv = GATConv(16, 32)
+    conv.__class__.jit_on_init = False
+
+    old_propagate = conv.__class__.propagate
+    old_edge_updater = conv.__class__.edge_updater
+
+    try:
+        conv.__class__.propagate = conv.propagate_jit
+        conv.__class__.collect = conv.collect_jit
+        conv.__class__.edge_updater = conv.edge_updater_jit
+        conv.__class__.edge_collect = conv.edge_collect_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 = conv(x, edge_index)
+        out2 = conv(x, adj.t())
+        assert torch.allclose(out1, out2)
+
+    finally:
+        conv.__class__.propagate = old_propagate
+        conv.__class__.edge_updater = old_edge_updater