Repository: pyg-team/pytorch_geometric
Pull request #8796 (2024-01-19 21:17:18)

Refactor `Sequential` module
Use the new functionality to create jinja templates introduced as part of the `MessagePassing` refactoring.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `12 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`effff69`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/effff6924646e9ad1172701b40d9ea036e62da01?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.42% compared to head [(`2e46098`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.39%.

> :exclamation: Current head 2e46098 differs from pull request most recent head 8000d03. Consider uploading reports for the commit 8000d03 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?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/sequential.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3NlcXVlbnRpYWwucHk=) | 80.35% | [11 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |
| [torch\_geometric/inspector.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2luc3BlY3Rvci5weQ==) | 96.96% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?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    #8796      +/-   ##
==========================================
- Coverage   89.42%   89.39%   -0.03%     
==========================================
  Files         479      479              
  Lines       31073    31082       +9     
==========================================
  Hits        27787    27787              
- Misses       3286     3295       +9     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8796?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/inspector.py b/torch_geometric/inspector.py
index ad9a6d734..3de3c6531 100644
--- a/torch_geometric/inspector.py
+++ b/torch_geometric/inspector.py
@@ -41,54 +41,11 @@ class Inspector:
 
     def eval_type(self, value: Any) -> Type:
         r"""Returns the type hint of a string."""
-        if isinstance(value, str):
-            value = typing.ForwardRef(value)
-        return typing._eval_type(value, self._globals, None)  # type: ignore
+        return eval_type(value, self._globals)
 
     def type_repr(self, obj: Any) -> str:
         r"""Returns the type hint representation of an object."""
-        def _get_name(name: str, module: str) -> str:
-            return name if name in self._globals else f'{module}.{name}'
-
-        if isinstance(obj, str):
-            return obj
-
-        if obj is type(None):
-            return 'None'
-
-        if obj is ...:
-            return '...'
-
-        if obj.__module__ == 'typing':  # Special logic for `typing.*` types:
-            name = obj._name
-            if name is None:  # In some cases, `_name` is not populated.
-                name = str(obj.__origin__).split('.')[-1]
-
-            args = getattr(obj, '__args__', None)
-            if args is None or len(args) == 0:
-                return _get_name(name, obj.__module__)
-            if all(isinstance(arg, typing.TypeVar) for arg in args):
-                return _get_name(name, obj.__module__)
-
-            # Convert `Union[*, None]` to `Optional[*]`.
-            # This is only necessary for old Python versions, e.g. 3.8.
-            # TODO Only convert to `Optional` if `Optional` is importable.
-            if (name == 'Union' and len(args) == 2
-                    and any([arg is type(None) for arg in args])
-                    and ('Optional' in self._globals
-                         or obj.__module__ in self._globals or True)):
-                name = 'Optional'
-
-            if name == 'Optional':  # Remove `None` from `Optional` arguments:
-                args = [arg for arg in obj.__args__ if arg is not type(None)]
-
-            args_repr = ', '.join([self.type_repr(arg) for arg in args])
-            return f'{_get_name(name, obj.__module__)}[{args_repr}]'
-
-        if obj.__module__ == 'builtins':
-            return obj.__qualname__
-
-        return _get_name(obj.__qualname__, obj.__module__)
+        return type_repr(obj, self._globals)
 
     def implements(self, func_name: str) -> bool:
         r"""Returns :obj:`True` in case the inspected class implements the
@@ -458,6 +415,57 @@ class Inspector:
         return {}  # (4) No function call found:
 
 
+def eval_type(value: Any, _globals: Dict[str, Any]) -> Type:
+    r"""Returns the type hint of a string."""
+    if isinstance(value, str):
+        value = typing.ForwardRef(value)
+    return typing._eval_type(value, _globals, None)  # type: ignore
+
+
+def type_repr(obj: Any, _globals: Dict[str, Any]) -> str:
+    r"""Returns the type hint representation of an object."""
+    def _get_name(name: str, module: str) -> str:
+        return name if name in _globals else f'{module}.{name}'
+
+    if isinstance(obj, str):
+        return obj
+
+    if obj is type(None):
+        return 'None'
+
+    if obj is ...:
+        return '...'
+
+    if obj.__module__ == 'typing':  # Special logic for `typing.*` types:
+        name = obj._name
+        if name is None:  # In some cases, `_name` is not populated.
+            name = str(obj.__origin__).split('.')[-1]
+
+        args = getattr(obj, '__args__', None)
+        if args is None or len(args) == 0:
+            return _get_name(name, obj.__module__)
+        if all(isinstance(arg, typing.TypeVar) for arg in args):
+            return _get_name(name, obj.__module__)
+
+        # Convert `Union[*, None]` to `Optional[*]`.
+        # This is only necessary for old Python versions, e.g. 3.8.
+        # TODO Only convert to `Optional` if `Optional` is importable.
+        if (name == 'Union' and len(args) == 2
+                and any([arg is type(None) for arg in args])):
+            name = 'Optional'
+
+        if name == 'Optional':  # Remove `None` from `Optional` arguments:
+            args = [arg for arg in obj.__args__ if arg is not type(None)]
+
+        args_repr = ', '.join([type_repr(arg, _globals) for arg in args])
+        return f'{_get_name(name, obj.__module__)}[{args_repr}]'
+
+    if obj.__module__ == 'builtins':
+        return obj.__qualname__
+
+    return _get_name(obj.__qualname__, obj.__module__)
+
+
 def find_parenthesis_content(source: str, prefix: str) -> Optional[str]:
     r"""Returns the content of :obj:`{prefix}.*(...)` within :obj:`source`."""
     match = re.search(prefix, source)
diff --git a/torch_geometric/nn/sequential.jinja b/torch_geometric/nn/sequential.jinja
index 16bbfaea6..a64713eb5 100644
--- a/torch_geometric/nn/sequential.jinja
+++ b/torch_geometric/nn/sequential.jinja
@@ -1,33 +1,35 @@
+import typing
 from typing import *
-from torch_geometric.typing import *
 
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
+from torch_geometric.typing import *
+
 
-class {{cls_name}}(torch.nn.Module):
-    def reset_parameters(self):
-        for name in self._names:
-            module = getattr(self, name)
-            if hasattr(module, 'reset_parameters'):
-                module.reset_parameters()
+class Sequential(torch.nn.Module):
+    def reset_parameters(self) -> None:
+{%- for child in children %}
+        if hasattr(self.{{child.name}}, 'reset_parameters'):
+            self.{{child.name}}.reset_parameters()
+{%- endfor %}
 
-    def forward(self, {{ input_args|join(', ') }}):
-        """"""{% for name, module, in_desc, out_desc in calls %}
-        {% for output in out_desc %}{{output}}{{ ", " if not loop.last }}{% endfor %} = self.{{name}}({% for input in in_desc %}{{input}}{{ ", " if not loop.last }}{% endfor %}){% endfor %}
-        return {% for output in calls[-1][-1] %}{{output}}{{ ", " if not loop.last }}{% endfor %}
+    def forward(self, {{ input_types|join(', ') }}) -> {{return_type}}:
+{%- for child in children %}
+        {{child.return_names|join(', ')}} = self.{{child.name}}({{child.param_names|join(', ')}})
+{%- endfor %}
+        return {{children[-1].return_names|join(', ')}}
 
     def __getitem__(self, idx: int) -> torch.nn.Module:
-        return getattr(self, self._names[idx])
+        return getattr(self, self._module_names[idx])
 
     def __len__(self) -> int:
-        return {{calls|length}}
+        return {{children|length}}
 
     def __repr__(self) -> str:
         module_reprs = [
-            (f"  ({i}) - {self[i]}: {', '.join(self.module_headers[i].args)} "
-             f"-> {', '.join(self.module_headers[i].output)}")
+            f'  ({i}) - {self[i]}: {self._module_descs[i]}'
             for i in range(len(self))
         ]
-
         return 'Sequential(\n{}\n)'.format('\n'.join(module_reprs))
diff --git a/torch_geometric/nn/sequential.py b/torch_geometric/nn/sequential.py
index 67900a0fb..498eaa78c 100644
--- a/torch_geometric/nn/sequential.py
+++ b/torch_geometric/nn/sequential.py
@@ -1,31 +1,19 @@
-import os
 import os.path as osp
-import sys
-from getpass import getuser
-from importlib.util import module_from_spec, spec_from_file_location
-from tempfile import NamedTemporaryFile as TempFile
-from tempfile import gettempdir
-from typing import Callable, List, NamedTuple, Tuple, Union
-from uuid import uuid1
+import random
+from typing import Callable, List, NamedTuple, Optional, Tuple, Union
 
 import torch
+from torch import Tensor
 
+from torch_geometric.inspector import split, type_repr
+from torch_geometric.template import module_from_template
 
-def class_from_module_repr(cls_name, module_repr):
-    path = osp.join(gettempdir(), f'{getuser()}_pyg')
-    os.makedirs(path, exist_ok=True)
-    with TempFile(mode='w+', suffix='.py', delete=False, dir=path) as f:
-        f.write(module_repr)
-    spec = spec_from_file_location(cls_name, f.name)
-    mod = module_from_spec(spec)
-    sys.modules[cls_name] = mod
-    spec.loader.exec_module(mod)
-    return getattr(mod, cls_name)
 
-
-class HeaderDesc(NamedTuple):
-    args: List[str]
-    output: List[str]
+class Child(NamedTuple):
+    name: str
+    module: Callable
+    param_names: List[str]
+    return_names: List[str]
 
 
 def Sequential(
@@ -36,8 +24,8 @@ def Sequential(
     define a sequential GNN model.
 
     Since GNN operators take in multiple input arguments,
-    :class:`torch_geometric.nn.Sequential` expects both global input
-    arguments, and function header definitions of individual operators.
+    :class:`torch_geometric.nn.Sequential` additionally expects both global
+    input arguments, and function header definitions of individual operators.
     If omitted, an intermediate module will operate on the *output* of its
     preceding module:
 
@@ -54,9 +42,9 @@ def Sequential(
             Linear(64, out_channels),
         ])
 
-    where :obj:`'x, edge_index'` defines the input arguments of :obj:`model`,
+    Here, :obj:`'x, edge_index'` defines the input arguments of :obj:`model`,
     and :obj:`'x, edge_index -> x'` defines the function header, *i.e.* input
-    arguments *and* return types, of :class:`~torch_geometric.nn.conv.GCNConv`.
+    arguments *and* return types of :class:`~torch_geometric.nn.conv.GCNConv`.
 
     In particular, this also allows to create more sophisticated models,
     such as utilizing :class:`~torch_geometric.nn.models.JumpingKnowledge`:
@@ -86,62 +74,74 @@ def Sequential(
             :obj:`OrderedDict` of modules (and function header definitions) can
             be passed.
     """
-    try:
-        from jinja2 import Template
-    except ImportError:
-        raise ModuleNotFoundError(
-            "No module named 'jinja2' found on this machine. "
-            "Run 'pip install jinja2' to install the library.")
-
-    input_args = [x.strip() for x in input_args.split(',')]
+    signature = input_args.split('->')
+    if len(signature) == 1:
+        input_args = signature[0]
+        return_type = type_repr(Tensor, globals())
+    elif len(signature) == 2:
+        input_args, return_type = signature[0], signature[1].strip()
+    else:
+        raise ValueError(f"Failed to parse arguments (got '{input_args}')")
+
+    input_types = split(input_args, sep=',')
+    if len(input_types) == 0:
+        raise ValueError(f"Failed to parse arguments (got '{input_args}')")
 
     if not isinstance(modules, dict):
         modules = {f'module_{i}': module for i, module in enumerate(modules)}
-
-    # We require the first entry of the input list to define arguments:
-    assert len(modules) > 0
-    first_module = list(modules.values())[0]
-    assert isinstance(first_module, (tuple, list))
-
-    # A list holding the callable function and the input and output names:
-    calls: List[Tuple[str, Callable, List[str], List[str]]] = []
-
-    for name, module in modules.items():
-        if isinstance(module, (tuple, list)) and len(module) >= 2:
-            module, desc = module[:2]
-            in_desc, out_desc = parse_desc(desc)
-        elif isinstance(module, (tuple, list)):
-            module = module[0]
-            in_desc = out_desc = calls[-1][-1]
+    if len(modules) == 0:
+        raise ValueError("'Sequential' expected a non-empty list of modules")
+
+    children: List[Child] = []
+    for i, (name, module) in enumerate(modules.items()):
+        desc: Optional[str] = None
+        if isinstance(module, (tuple, list)):
+            if len(module) == 1:
+                module = module[0]
+            elif len(module) == 2:
+                module, desc = module
+            else:
+                raise ValueError(f"Expected tuple of length 2 (got {module})")
+
+        if i == 0 and desc is None:
+            raise ValueError("Requires signature for first module")
+        if not callable(module):
+            raise ValueError(f"Expected callable module (got {module})")
+        if desc is not None and not isinstance(desc, str):
+            raise ValueError(f"Expected type hint representation (got {desc})")
+
+        if desc is not None:
+            signature = desc.split('->')
+            if len(signature) != 2:
+                raise ValueError(f"Failed to parse arguments (got '{desc}')")
+            param_names = [v.strip() for v in signature[0].split(',')]
+            return_names = [v.strip() for v in signature[1].split(',')]
+            child = Child(name, module, param_names, return_names)
         else:
-            in_desc = out_desc = calls[-1][-1]
-
-        calls.append((name, module, in_desc, out_desc))
-
-    root = osp.dirname(osp.realpath(__file__))
-    with open(osp.join(root, 'sequential.jinja'), 'r') as f:
-        template = Template(f.read())
-
-    cls_name = f'Sequential_{uuid1().hex[:6]}'
-    module_repr = template.render(
-        cls_name=cls_name,
-        input_args=input_args,
-        calls=calls,
+            param_names = children[-1].return_names
+            child = Child(name, module, param_names, param_names)
+
+        children.append(child)
+
+    uid = '%06x' % random.randrange(16**6)
+    root_dir = osp.dirname(osp.realpath(__file__))
+    module = module_from_template(
+        module_name=f'torch_geometric.nn.sequential_{uid}',
+        template_path=osp.join(root_dir, 'sequential.jinja'),
+        tmp_dirname='sequential',
+        # Keyword arguments:
+        input_types=input_types,
+        return_type=return_type,
+        children=children,
     )
 
-    # Instantiate a class from the rendered module representation.
-    module = class_from_module_repr(cls_name, module_repr)()
-    module.module_headers = [
-        HeaderDesc(in_desc, out_desc) for _, _, in_desc, out_desc in calls
+    model = module.Sequential()
+    model._module_names = [child.name for child in children]
+    model._module_descs = [
+        f"{', '.join(child.param_names)} -> {', '.join(child.return_names)}"
+        for child in children
     ]
-    module._names = list(modules.keys())
-    for name, submodule, _, _ in calls:
-        setattr(module, name, submodule)
-    return module
-
+    for child in children:
+        setattr(model, child.name, child.module)
 
-def parse_desc(desc: str) -> Tuple[List[str], List[str]]:
-    in_desc, out_desc = desc.split('->')
-    in_desc = [x.strip() for x in in_desc.split(',')]
-    out_desc = [x.strip() for x in out_desc.split(',')]
-    return in_desc, out_desc
+    return model


Test patch:
diff --git a/test/nn/test_sequential.py b/test/nn/test_sequential.py
index 6f61fab47..2c997af0a 100644
--- a/test/nn/test_sequential.py
+++ b/test/nn/test_sequential.py
@@ -47,11 +47,11 @@ def test_sequential():
     assert isinstance(model[3], ReLU)
     assert isinstance(model[4], Linear)
 
-    assert model.module_headers[0] == (['x', 'edge_index'], ['x'])
-    assert model.module_headers[1] == (['x'], ['x'])
-    assert model.module_headers[2] == (['x', 'edge_index'], ['x'])
-    assert model.module_headers[3] == (['x'], ['x'])
-    assert model.module_headers[4] == (['x'], ['x'])
+    assert model._module_descs[0] == 'x, edge_index -> x'
+    assert model._module_descs[1] == 'x -> x'
+    assert model._module_descs[2] == 'x, edge_index -> x'
+    assert model._module_descs[3] == 'x -> x'
+    assert model._module_descs[4] == 'x -> x'
 
     out = model(x, edge_index)
     assert out.size() == (4, 7)
@@ -73,14 +73,14 @@ def test_sequential():
     assert out.size() == (1, 7)
 
 
-def test_sequential_jittable():
+def test_sequential_jit():
     x = torch.randn(4, 16)
     edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
 
     model = Sequential('x: Tensor, edge_index: Tensor', [
-        (GCNConv(16, 64).jittable(), 'x, edge_index -> x'),
+        (GCNConv(16, 64), 'x, edge_index -> x'),
         ReLU(inplace=True),
-        (GCNConv(64, 64).jittable(), 'x, edge_index -> x'),
+        (GCNConv(64, 64), 'x, edge_index -> x'),
         ReLU(inplace=True),
         Linear(64, 7),
     ])
@@ -90,9 +90,9 @@ def test_sequential_jittable():
         adj_t = SparseTensor.from_edge_index(edge_index).t()
 
         model = Sequential('x: Tensor, edge_index: SparseTensor', [
-            (GCNConv(16, 64).jittable(), 'x, edge_index -> x'),
+            (GCNConv(16, 64), 'x, edge_index -> x'),
             ReLU(inplace=True),
-            (GCNConv(64, 64).jittable(), 'x, edge_index -> x'),
+            (GCNConv(64, 64), 'x, edge_index -> x'),
             ReLU(inplace=True),
             Linear(64, 7),
         ])