Repository: pyg-team/pytorch_geometric
Pull request #8745 (2024-01-10 06:16:19)

`[MessagePassing]` Investigate `jittable` conversion on instantiation of GNN layers

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

> :exclamation: Current head 696ce6d differs from pull request most recent head 6da7591. Consider uploading reports for the commit 6da7591 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8745?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/nn/conv/message\_passing.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8745?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvbWVzc2FnZV9wYXNzaW5nLnB5) | 88.00% | [3 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8745?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |
| [torch\_geometric/nn/conv/propagate.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8745?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvcHJvcGFnYXRlLnB5) | 95.65% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8745?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    #8745      +/-   ##
==========================================
- Coverage   90.08%   89.56%   -0.53%     
==========================================
  Files         481      482       +1     
  Lines       31060    31112      +52     
==========================================
- Hits        27979    27864     -115     
- Misses       3081     3248     +167     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8745?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/CHANGELOG.md b/CHANGELOG.md
index 4e71ac0e5..ba9ff52eb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- GNN layers are now jittable by default ([#8745](https://github.com/pyg-team/pytorch_geometric/pull/8745))
 - Sparse node features in `NELL` and `AttributedGraphDataset` are now represented as `torch.sparse_csr_tensor` instead of `torch_sparse.SparseTensor` ([#8679](https://github.com/pyg-team/pytorch_geometric/pull/8679))
 - Accelerated mini-batching of `torch.sparse` tensors ([#8670](https://github.com/pyg-team/pytorch_geometric/pull/8670))
 - Fixed RPC timeout due to worker closing in `DistLoader` with `atexit` not executed correctly in `worker_init_fn` ([#8605](https://github.com/pyg-team/pytorch_geometric/pull/8605))
diff --git a/examples/README.md b/examples/README.md
new file mode 100644
index 000000000..25fe7a7cb
--- /dev/null
+++ b/examples/README.md
@@ -0,0 +1,20 @@
+# Examples
+
+This folder contains a plethora of examples covering different GNN use-cases.
+This readme highlights some key examples.
+
+A great and simple example to start with is [`gcn.py`](./gcn.py), showing a user how to train a [`GCN`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.GCN.html) model for node-level prediction on small-scale homogeneous data.
+
+For a simple link prediction example, see [`link_pred.py`](./link_pred.py).
+
+For examples on [Open Graph Benchmark](https://ogb.stanford.edu/) datasets, see the `ogbn_*.py` examples:
+
+- [`ogbn_products_sage.py`](./ogbn_products_sage.py) and [`ogbn_products_gat.py`](./ogbn_products_gat.py) show how to train [`GraphSAGE`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.GraphSAGE.html) and [`GAT`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.GAT.html) models on the `ogbn-products` dataset.
+- [`ogbn_proteins_deepgcn.py`](./ogbn_proteins_deepgcn.py) is an example to showcase how to train deep GNNs on the `ogbn-proteins` dataset.
+- [`ogbn_papers_100m.py`](./ogbn_papers_100m.py) is an example for training a GNN on the large-scale `ogbn-papers100m` dataset, containing approximately ~1.6B edges.
+
+For examples on using `torch.compile`, see the examples under [`examples/compile`](./compile).
+
+For examples on scaling PyG up via multi-GPUs, see the examples under [`examples/multi_gpu`](./multi_gpu).
+
+For examples on working with heterogeneous data, see the examples under [`examples/hetero`](./hetero).
diff --git a/examples/hetero/README.md b/examples/hetero/README.md
new file mode 100644
index 000000000..55101e1d1
--- /dev/null
+++ b/examples/hetero/README.md
@@ -0,0 +1,16 @@
+# Examples for Heterogeneous Data
+
+| Example                                                | Description                                                                                                                            |
+| ------------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------- |
+| [`hetero_conv_dblp.py`](./hetero_conv_dblp.py)         | Shows how to use the `HeteroConv(...)` wrapper; Trains it for node classification on the `DBLP` dataset.                               |
+| [`to_hetero_mag.py`](./to_hetero_mag.py)               | Shows how to use the `to_hetero(...)` functionality; Trains it for node classification on the `ogb-mag` dataset.                       |
+| [`hetero_link_pred.py`](./hetero_link_pred.py)         | Shows how to use the `to_hetero(...)` functionality; Trains it for link prediction on the `MovieLens` dataset.                         |
+| [`hgt_dblp.py`](./hgt_dblp.py)                         | Trains a Heterogeneous Graph Transformer (HGT) model for node classification on the `DBLP` dataset.                                    |
+| [`hierarchical_sage.py`](./hierarchical_sage.py)       | Shows how to perform hierarchical sampling; Trains a heterogeneous `GraphSAGE` model for node classification on the `ogb-mag` dataset. |
+| [`load_csv.py`](./load_csv.py)                         | Shows how to create heterogeneous graphs from raw `*.csv` data.                                                                        |
+| [`metapath2vec.py`](./metapath2vec.py)                 | Train an unsupervised `MetaPath2Vec` model; Tests embeddings for node classification on the `AMiner` dataset.                          |
+| [`temporal_link_pred.py`](./temporal_link_pred.py)     | Trains a heterogeneous `GraphSAGE` model for temporal link prediction on the `MovieLens` dataset.                                      |
+| [`bipartite_sage.py`](./bipartite_sage.py)             | Trains a GNN via metapaths for link prediction on the `MovieLens` dataset.                                                             |
+| [`bipartite_sage_unsup.py`](./bipartite_sage_unsup.py) | Trains a GNN via metapaths for link prediction on the large-scale `TaoBao` dataset.                                                    |
+| [`dmgi_unsup.py`](./dmgi_unsup.py)                     | Shows how to learn embeddings on the `IMDB` dataset using the `DMGI` model.                                                            |
+| [`han_imdb.py`](./han_imdb.py)                         | Shows how to train a Heterogenous Graph Attention Network (HAN) for node classification on the `IMDB` dataset.                         |
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index b5c83bf5f..1199c5b0f 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1678,7 +1678,8 @@ def matmul(
 
     rowptr: Optional[Tensor] = None
     if out.layout == torch.sparse_csr:
-        rowptr, col = out.crow_indices(), out.col_indices()
+        rowptr = out.crow_indices().to(input.dtype)
+        col = out.col_indices().to(input.dtype)
         edge_index = torch._convert_indices_from_csr_to_coo(
             rowptr, col, out_int32=rowptr.dtype != torch.int64)
 
diff --git a/torch_geometric/home.py b/torch_geometric/home.py
index db939d9a3..ec1704063 100644
--- a/torch_geometric/home.py
+++ b/torch_geometric/home.py
@@ -17,9 +17,7 @@ def get_home_dir() -> str:
     if _home_dir is not None:
         return _home_dir
 
-    home_dir = os.getenv(ENV_PYG_HOME, DEFAULT_CACHE_DIR)
-    home_dir = osp.expanduser(home_dir)
-    return home_dir
+    return osp.expanduser(os.getenv(ENV_PYG_HOME, DEFAULT_CACHE_DIR))
 
 
 def set_home_dir(path: str) -> None:
diff --git a/torch_geometric/nn/conv/fa_conv.py b/torch_geometric/nn/conv/fa_conv.py
index 87346219b..b844ec982 100644
--- a/torch_geometric/nn/conv/fa_conv.py
+++ b/torch_geometric/nn/conv/fa_conv.py
@@ -201,7 +201,8 @@ class FAConv(MessagePassing):
         alpha_l = self.att_l(x)
         alpha_r = self.att_r(x)
 
-        # propagate_type: (x: Tensor, alpha: PairTensor, edge_weight: OptTensor)  # noqa
+        # propagate_type: (x: Tensor, alpha: PairTensor,
+        #                  edge_weight: OptTensor)
         out = self.propagate(edge_index, x=x, alpha=(alpha_l, alpha_r),
                              edge_weight=edge_weight)
 
diff --git a/torch_geometric/nn/conv/heat_conv.py b/torch_geometric/nn/conv/heat_conv.py
index be89ca9b7..e3c1a8426 100644
--- a/torch_geometric/nn/conv/heat_conv.py
+++ b/torch_geometric/nn/conv/heat_conv.py
@@ -107,7 +107,8 @@ class HEATConv(MessagePassing):
         edge_type_emb = F.leaky_relu(self.edge_type_emb(edge_type),
                                      self.negative_slope)
 
-        # propagate_type: (x: Tensor, edge_type_emb: Tensor, edge_attr: OptTensor)  # noqa
+        # propagate_type: (x: Tensor, edge_type_emb: Tensor,
+        #                  edge_attr: OptTensor)
         out = self.propagate(edge_index, x=x, edge_type_emb=edge_type_emb,
                              edge_attr=edge_attr)
 
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index d0045655e..578ef991b 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -22,6 +22,11 @@ from torch import Tensor
 from torch.utils.hooks import RemovableHandle
 
 from torch_geometric.nn.aggr import Aggregation
+from torch_geometric.nn.conv.propagate import (
+    find_parenthesis_content,
+    module_from_template,
+    type_hint_to_str,
+)
 from torch_geometric.nn.conv.utils.inspector import Inspector
 from torch_geometric.nn.conv.utils.jit import class_from_module_repr
 from torch_geometric.nn.conv.utils.typing import sanitize, split_types_repr
@@ -173,6 +178,60 @@ class MessagePassing(torch.nn.Module):
         self._edge_update_forward_pre_hooks: HookDict = OrderedDict()
         self._edge_update_forward_hooks: HookDict = OrderedDict()
 
+        # Test code for performing on-the-fly TorchScript support:
+        if getattr(self, 'jit_on_init', False):
+            root_dir = osp.dirname(osp.realpath(__file__))
+            prop_types, prop_return_type = self._get_propagate_types()
+            module = module_from_template(
+                module_name=f'{self.__module__}_propagate',
+                module=self.__module__,
+                template_path=osp.join(root_dir, 'propagate.jinja'),
+                propagate_types=prop_types,
+                propagate_return_type=prop_return_type,
+                collect_types=self.inspector.types(
+                    ['message', 'aggregate', 'update']),
+                message_args=self.inspector.keys(['message']),
+                aggregate_args=self.inspector.keys(['aggregate']),
+                message_and_aggregate_args=self.inspector.keys(
+                    ['message_and_aggregate']),
+                update_args=self.inspector.keys(['update']),
+            )
+
+            self.propagate_jit = module.propagate
+            self._collect_jit = module._collect
+
+    def _get_propagate_types(self) -> Tuple[Dict[str, str], str]:
+        # Parse `propagate_types` and generate an efficient `propagate` method:
+        if hasattr(self, 'propagate_type'):
+            assert isinstance(self.propagate_type, dict)
+            propagate_types = {
+                name: type_hint_to_str(type_hint)
+                for name, type_hint in self.propagate_type.items()
+            }
+        else:
+            source = inspect.getsource(self.forward)
+            match = find_parenthesis_content(source, prefix='propagate_type:')
+            if match is not None:
+                propagate_types = dict(
+                    [re.split(r'\s*:\s*', t) for t in split_types_repr(match)])
+            else:
+                match = find_parenthesis_content(source, 'self.propagate')
+                if match is None:  # No `self.propagate` call:
+                    return
+                args = [  # Find all keyword argument names:
+                    arg[:arg.find('=')].strip() for arg in match.split(',')
+                    if arg.find('=') >= 0
+                ]
+                propagate_types = {
+                    arg: 'Tensor'
+                    for arg in args if arg not in {'edge_index', 'size'}
+                }
+
+        propagate_return_type = type_hint_to_str(
+            get_type_hints(self.update).get('return', Tensor))
+
+        return propagate_types, propagate_return_type
+
     def reset_parameters(self) -> None:
         r"""Resets all learnable parameters of the module."""
         if self.aggr_module is not None:
@@ -845,26 +904,7 @@ class MessagePassing(torch.nn.Module):
         source = inspect.getsource(self.__class__)
 
         # Find and parse `propagate()` types to format `{arg1: type1, ...}`.
-        if hasattr(self, 'propagate_type'):
-            assert isinstance(self.propagate_type, dict)
-            prop_types = {
-                k: sanitize(str(v))
-                for k, v in self.propagate_type.items()
-            }
-        else:
-            match = re.search(r'#\s*propagate_type:\s*\((.*)\)', source)
-            if match is None:
-                raise TypeError(
-                    'TorchScript support requires the definition of the types '
-                    'passed to `propagate()`. Please specify them via\n\n'
-                    'propagate_type = {"arg1": type1, "arg2": type2, ... }\n\n'
-                    'or via\n\n'
-                    '# propagate_type: (arg1: type1, arg2: type2, ...)\n\n'
-                    'inside the `MessagePassing` module.')
-            prop_types = dict([
-                re.split(r'\s*:\s*', t)
-                for t in split_types_repr(match.group(1))
-            ])
+        prop_types, prop_return_type = self._get_propagate_types()
 
         # Find and parse `edge_updater` types to format `{arg1: type1, ...}`.
         if 'edge_update' in self.__class__.__dict__.keys():
@@ -875,7 +915,7 @@ class MessagePassing(torch.nn.Module):
                     for k, v in self.edge_updater_type.items()
                 }
             else:
-                match = re.search(r'#\s*edge_updater_type:\s*\((.*)\)', source)
+                match = find_parenthesis_content(source, 'edge_updater_type:')
                 if match is None:
                     raise TypeError(
                         'TorchScript support requires the definition of the '
@@ -884,18 +924,11 @@ class MessagePassing(torch.nn.Module):
                         '"arg2": type2, ... }\n\n or via\n\n'
                         '# edge_updater_type: (arg1: type1, arg2: type2, ...)'
                         '\n\ninside the `MessagePassing` module.')
-                edge_updater_types = dict([
-                    re.split(r'\s*:\s*', t)
-                    for t in split_types_repr(match.group(1))
-                ])
+                edge_updater_types = dict(
+                    [re.split(r'\s*:\s*', t) for t in split_types_repr(match)])
         else:
             edge_updater_types = {}
 
-        type_hints = get_type_hints(self.__class__.update)
-        prop_return_type = type_hints.get('return', 'Tensor')
-        if str(prop_return_type)[:6] == '<class':
-            prop_return_type = prop_return_type.__name__
-
         type_hints = get_type_hints(self.__class__.edge_update)
         edge_updater_return_type = type_hints.get('return', 'Tensor')
         if str(edge_updater_return_type)[:6] == '<class':
diff --git a/torch_geometric/nn/conv/ppf_conv.py b/torch_geometric/nn/conv/ppf_conv.py
index 5ca2547d5..0463e2213 100644
--- a/torch_geometric/nn/conv/ppf_conv.py
+++ b/torch_geometric/nn/conv/ppf_conv.py
@@ -121,7 +121,8 @@ class PPFConv(MessagePassing):
             elif isinstance(edge_index, SparseTensor):
                 edge_index = torch_sparse.set_diag(edge_index)
 
-        # propagate_type: (x: PairOptTensor, pos: PairTensor, normal: PairTensor)  # noqa
+        # propagate_type: (x: PairOptTensor, pos: PairTensor,
+        #                  normal: PairTensor)
         out = self.propagate(edge_index, x=x, pos=pos, normal=normal)
 
         if self.global_nn is not None:
diff --git a/torch_geometric/nn/conv/propagate.jinja b/torch_geometric/nn/conv/propagate.jinja
new file mode 100644
index 000000000..9b6e46017
--- /dev/null
+++ b/torch_geometric/nn/conv/propagate.jinja
@@ -0,0 +1,105 @@
+import typing
+from typing import *
+
+import torch
+from torch import Tensor
+
+from torch_geometric.typing import *
+from torch_geometric import is_compiling
+from torch_geometric.utils import is_sparse
+
+from {{module}} import *
+
+
+class CollectArgs(NamedTuple):
+{%- for name, type_hint in collect_types.items() %}
+    {{name}}: {{type_hint}}
+{%- endfor %}
+
+
+def _collect(
+    self,
+    edge_index: Union[Tensor, SparseTensor],
+{%- for name, type_hint in propagate_types.items() %}
+    {{name}}: {{type_hint}},
+{%- endfor %}
+    size: List[Optional[int]],
+) -> CollectArgs:
+
+    i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
+{% for name in collect_types %}
+{%- if name.endswith('_i') or name.endswith('_j') %}
+    # 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]}})
+        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]}})
+        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]}})
+    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
+
+    return CollectArgs({% for name in collect_types %}{{name}}{% if not loop.last %}, {% endif %}{% endfor %})
+
+
+def propagate(
+    self,
+    edge_index: Union[Tensor, SparseTensor],
+{%- for name, type_hint in propagate_types.items() %}
+    {{name}}: {{type_hint}},
+{%- endfor %}
+    size: Size = None,
+) -> {{propagate_return_type}}:
+
+    decomposed_layers = 1 if self.explain is True else self.decomposed_layers
+
+    if not torch.jit.is_scripting() and not is_compiling():
+        for hook in self._propagate_forward_pre_hooks.values():
+            kwargs = dict({% for name in propagate_types %}{{name}}={{name}}{% if not loop.last %}, {% endif %}{% endfor %})
+            res = hook(self, (edge_index, size, kwargs))
+            if res is not None:
+                edge_index, size, kwargs = res
+{%- for name, type_hint in propagate_types.items() %}
+                {{name}} = kwargs['{{name}}']
+{%- endfor %}
+
+    mutable_size = self._check_input(edge_index, size)
+    fuse = is_sparse(edge_index) and self.fuse and self.explain is not True
+
+    if fuse:
+        raise NotImplementedError
+
+    else:
+        kwargs = self._collect(edge_index, {% for name in propagate_types %}{{name}}, {% endfor %}mutable_size)
+        out = self.message({% for name in message_args %}{{name}}=kwargs.{{name}}{% if not loop.last %}, {% endif %}{% endfor %})
+        out = self.aggregate(out{% for name in aggregate_args %}, {{name}}=kwargs.{{name}}{% endfor %})
+        out = self.update(out{% for name in update_args %}, {{name}}=kwargs.{{name}}{% endfor %})
+        return out
diff --git a/torch_geometric/nn/conv/propagate.py b/torch_geometric/nn/conv/propagate.py
new file mode 100644
index 000000000..2dc629e71
--- /dev/null
+++ b/torch_geometric/nn/conv/propagate.py
@@ -0,0 +1,65 @@
+import importlib
+import os
+import os.path as osp
+import re
+import sys
+from typing import Any, Optional
+
+from jinja2 import Template
+
+from torch_geometric import get_home_dir
+
+
+def module_from_template(
+    module_name: str,
+    template_path: str,
+    **kwargs: Any,
+) -> Any:
+
+    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())
+    module_repr = template.render(**kwargs)
+
+    instance_dir = osp.join(get_home_dir(), 'propagate')
+    os.makedirs(instance_dir, exist_ok=True)
+    instance_path = osp.join(instance_dir, f'{module_name}.py')
+    with open(instance_path, 'w') as f:
+        f.write(module_repr)
+
+    spec = importlib.util.spec_from_file_location(module_name, instance_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def type_hint_to_str(type_hint: Any) -> str:
+    type_repr = str(type_hint)
+    type_repr = re.sub(r'<class \'(.*)\'>', r'\1', type_repr)
+    return type_repr
+
+
+def find_parenthesis_content(text: str, prefix: str) -> Optional[str]:
+    match = re.search(prefix, text)
+    if match is None:
+        return
+
+    offset = text[match.start():].find('(')
+    if offset < 0:
+        return
+
+    content = text[match.start() + offset:]
+
+    num_opened = num_closed = 0
+    for end, char in enumerate(content):
+        if char == '(':
+            num_opened += 1
+        if char == ')':
+            num_closed += 1
+        if num_opened > 0 and num_opened == num_closed:
+            content = content[1:end]
+            content = content.replace('\n', ' ').replace('#', ' ')
+            return re.sub(' +', ' ', content.replace('\n', ' ')).strip()
diff --git a/torch_geometric/nn/conv/res_gated_graph_conv.py b/torch_geometric/nn/conv/res_gated_graph_conv.py
index 960d65b11..023d0920b 100644
--- a/torch_geometric/nn/conv/res_gated_graph_conv.py
+++ b/torch_geometric/nn/conv/res_gated_graph_conv.py
@@ -123,7 +123,8 @@ class ResGatedGraphConv(MessagePassing):
         else:
             k, q, v = x[1], x[0], x[0]
 
-        # propagate_type: (k: Tensor, q: Tensor, v: Tensor, edge_attr: OptTensor)  # noqa
+        # propagate_type: (k: Tensor, q: Tensor, v: Tensor,
+        #                  edge_attr: OptTensor)
         out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr)
 
         if self.root_weight:
diff --git a/torch_geometric/nn/conv/rgat_conv.py b/torch_geometric/nn/conv/rgat_conv.py
index 648d0380e..4c3a9937d 100644
--- a/torch_geometric/nn/conv/rgat_conv.py
+++ b/torch_geometric/nn/conv/rgat_conv.py
@@ -345,7 +345,8 @@ class RGATConv(MessagePassing):
                 :obj:`(edge_index, attention_weights)`, holding the computed
                 attention weights for each edge. (default: :obj:`None`)
         """
-        # propagate_type: (x: Tensor, edge_type: OptTensor, edge_attr: OptTensor)  # noqa
+        # propagate_type: (x: Tensor, edge_type: OptTensor,
+        #                  edge_attr: OptTensor)
         out = self.propagate(edge_index=edge_index, edge_type=edge_type, x=x,
                              size=size, edge_attr=edge_attr)
 
diff --git a/torch_geometric/nn/conv/transformer_conv.py b/torch_geometric/nn/conv/transformer_conv.py
index 049addd6b..f6710efea 100644
--- a/torch_geometric/nn/conv/transformer_conv.py
+++ b/torch_geometric/nn/conv/transformer_conv.py
@@ -224,7 +224,8 @@ class TransformerConv(MessagePassing):
         key = self.lin_key(x[0]).view(-1, H, C)
         value = self.lin_value(x[0]).view(-1, H, C)
 
-        # propagate_type: (query: Tensor, key:Tensor, value: Tensor, edge_attr: OptTensor) # noqa
+        # propagate_type: (query: Tensor, key:Tensor, value: Tensor,
+        #                  edge_attr: OptTensor)
         out = self.propagate(edge_index, query=query, key=key, value=value,
                              edge_attr=edge_attr)
 


Test patch:
diff --git a/test/nn/conv/test_propagate.py b/test/nn/conv/test_propagate.py
new file mode 100644
index 000000000..ec662c7f3
--- /dev/null
+++ b/test/nn/conv/test_propagate.py
@@ -0,0 +1,17 @@
+import torch
+
+from torch_geometric.nn import SAGEConv
+
+
+def test_propagate():
+    SAGEConv.jit_on_init = True
+    conv = SAGEConv(10, 16)
+    SAGEConv.jit_on_init = False
+
+    old_propagate = conv.__class__.propagate
+    old_collect = conv.__class__._collect
+    conv.__class__.propagate = conv.propagate_jit
+    conv.__class__._collect = conv._collect_jit
+    torch.jit.script(conv)
+    conv.__class__.propagate = old_propagate
+    conv.__class__._collect = old_collect