Repository: pyg-team/pytorch_geometric
Pull request #8093 (2023-09-29 09:33:56)

Add `module_headers` property to `nn.Sequential` models
Fixes https://github.com/pyg-team/pytorch_geometric/issues/8082.

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

> :exclamation: Current head e879f02 differs from pull request most recent head fca34ff. Consider uploading reports for the commit fca34ff to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #8093   +/-   ##
=======================================
  Coverage   88.72%   88.72%           
=======================================
  Files         467      467           
  Lines       27897    27901    +4     
=======================================
+ Hits        24751    24755    +4     
  Misses       3146     3146           
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8093?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/sequential.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8093?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3NlcXVlbnRpYWwucHk=) | `91.30% <100.00%> (+0.82%)` | :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 c54a41c00..190cceba0 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 `module_headers` property to `nn.Sequential` models ([#8093](https://github.com/pyg-team/pytorch_geometric/pull/8093)
 - Added `OnDiskDataset` interface with data loader support ([#8066](https://github.com/pyg-team/pytorch_geometric/pull/8066), [#8088](https://github.com/pyg-team/pytorch_geometric/pull/8088), [#8092](https://github.com/pyg-team/pytorch_geometric/pull/8092))
 - Added a tutorial for `Node2Vec` and `MetaPath2Vec` usage ([#7938](https://github.com/pyg-team/pytorch_geometric/pull/7938)
 - Added a tutorial for multi-GPU training with pure PyTorch ([#7894](https://github.com/pyg-team/pytorch_geometric/pull/7894)
diff --git a/torch_geometric/nn/sequential.jinja b/torch_geometric/nn/sequential.jinja
index 013d6f46c..16bbfaea6 100644
--- a/torch_geometric/nn/sequential.jinja
+++ b/torch_geometric/nn/sequential.jinja
@@ -24,5 +24,10 @@ class {{cls_name}}(torch.nn.Module):
         return {{calls|length}}
 
     def __repr__(self) -> str:
-        return 'Sequential(\n{}\n)'.format('\n'.join(
-            [f'  ({idx}): ' + str(self[idx]) for idx in range(len(self))]))
+        module_reprs = [
+            (f"  ({i}) - {self[i]}: {', '.join(self.module_headers[i].args)} "
+             f"-> {', '.join(self.module_headers[i].output)}")
+            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 af4d4a88f..3a2ea7e01 100644
--- a/torch_geometric/nn/sequential.py
+++ b/torch_geometric/nn/sequential.py
@@ -1,6 +1,6 @@
 import os
 import os.path as osp
-from typing import Callable, List, Tuple, Union
+from typing import Callable, List, NamedTuple, Tuple, Union
 from uuid import uuid1
 
 import torch
@@ -8,6 +8,11 @@ import torch
 from torch_geometric.nn.conv.utils.jit import class_from_module_repr
 
 
+class HeaderDesc(NamedTuple):
+    args: List[str]
+    output: List[str]
+
+
 def Sequential(
     input_args: str,
     modules: List[Union[Tuple[Callable, str], Callable]],
@@ -110,6 +115,9 @@ def Sequential(
 
     # 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
+    ]
     module._names = list(modules.keys())
     for name, submodule, _, _ in calls:
         setattr(module, name, submodule)


Test patch:
diff --git a/test/nn/test_sequential.py b/test/nn/test_sequential.py
index 43adf1b27..6f61fab47 100644
--- a/test/nn/test_sequential.py
+++ b/test/nn/test_sequential.py
@@ -34,11 +34,11 @@ def test_sequential():
     assert len(model) == 5
     assert str(model) == (
         'Sequential(\n'
-        '  (0): GCNConv(16, 64)\n'
-        '  (1): ReLU(inplace=True)\n'
-        '  (2): GCNConv(64, 64)\n'
-        '  (3): ReLU(inplace=True)\n'
-        '  (4): Linear(in_features=64, out_features=7, bias=True)\n'
+        '  (0) - GCNConv(16, 64): x, edge_index -> x\n'
+        '  (1) - ReLU(inplace=True): x -> x\n'
+        '  (2) - GCNConv(64, 64): x, edge_index -> x\n'
+        '  (3) - ReLU(inplace=True): x -> x\n'
+        '  (4) - Linear(in_features=64, out_features=7, bias=True): x -> x\n'
         ')')
 
     assert isinstance(model[0], GCNConv)
@@ -47,6 +47,12 @@ 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'])
+
     out = model(x, edge_index)
     assert out.size() == (4, 7)