update

test/nn/conv/test_graph_conv.py

@@ -93,3 +93,12 @@ def test_graph_conv():
             assert torch.allclose(jit((x1, None), adj3.t()), out2, atol=1e-6)
             assert torch.allclose(jit((x1, x2), adj4.t()), out3, atol=1e-6)
             assert torch.allclose(jit((x1, None), adj4.t()), out4, atol=1e-6)
+
+
+class EdgeGraphConv(GraphConv):
+    def message(self, x_j, edge_weight):
+        return x_j if edge_weight is None else edge_weight * x_j
+
+
+def test_inheritance():
+    conv = EdgeGraphConv(16, 16)

torch_geometric/inspector.py

@@ -4,6 +4,7 @@
 import typing
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Type, Union
 
+import torch
 from torch import Tensor
 
 
@@ -32,9 +33,27 @@ def __init__(self, cls: Type):
         self._signature_dict: Dict[str, Signature] = {}
         self._source_dict: Dict[str, str] = {}
 
+    def _get_modules(self, cls: Type) -> List[str]:
+        from torch_geometric.nn import MessagePassing
+
+        modules: List[str] = []
+        for base_cls in cls.__bases__:
+            if base_cls not in {object, torch.nn.Module, MessagePassing}:
+                modules.extend(self._get_modules(base_cls))
+
+        modules.append(cls.__module__)
+        return modules
+
+    @property
+    def _modules(self) -> List[str]:
+        return self._get_modules(self._cls)
+
     @property
     def _globals(self) -> Dict[str, Any]:
-        return sys.modules[self._cls.__module__].__dict__
+        out: Dict[str, Any] = {}
+        for module in self._modules:
+            out.update(sys.modules[module].__dict__)
+        return out
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}({self._cls.__name__})'

torch_geometric/nn/conv/edge_updater.jinja

@@ -8,8 +8,9 @@ import torch_geometric.typing
 from torch_geometric import is_compiling
 from torch_geometric.utils import is_sparse
 from torch_geometric.typing import Size, SparseTensor
-
+{% for module in modules %}
 from {{module}} import *
+{%- endfor %}
 
 
 {% include "collect.jinja" %}

torch_geometric/nn/conv/message_passing.py

@@ -172,7 +172,7 @@ def __init__(
                     template_path=osp.join(root_dir, 'propagate.jinja'),
                     tmp_dirname='message_passing',
                     # Keyword arguments:
-                    module=self.__module__,
+                    modules=self.inspector._modules,
                     collect_name='collect',
                     signature=self._get_propagate_signature(),
                     collect_param_dict=self.inspector.get_flat_param_dict(
@@ -204,7 +204,7 @@ def __init__(
                 template_path=osp.join(root_dir, 'edge_updater.jinja'),
                 tmp_dirname='message_passing',
                 # Keyword arguments:
-                module=self.__module__,
+                modules=self.inspector._modules,
                 collect_name='edge_collect',
                 signature=self._get_edge_updater_signature(),
                 collect_param_dict=self.inspector.get_param_dict(

torch_geometric/nn/conv/propagate.jinja

@@ -8,8 +8,9 @@ import torch_geometric.typing
 from torch_geometric import is_compiling
 from torch_geometric.utils import is_sparse
 from torch_geometric.typing import Size, SparseTensor
-
+{% for module in modules %}
 from {{module}} import *
+{%- endfor %}
 
 
 {% include "collect.jinja" %}

torch_geometric/template.py

@@ -22,6 +22,7 @@ def module_from_template(
     env = Environment(loader=FileSystemLoader(osp.dirname(template_path)))
     template = env.get_template(osp.basename(template_path))
     module_repr = template.render(**kwargs)
+    print(module_repr)
 
     instance_dir = osp.join(get_home_dir(), tmp_dirname)
     os.makedirs(instance_dir, exist_ok=True)