Fix for configmixin with explicit classes (#10129)

Fix for configmixin with explicit classes

Author: fedelopez77

CHANGELOG.md

@@ -51,7 +51,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
-- Fixed `_recursive_config()` for `torch.nn.ModuleList`([#10124](https://github.com/pyg-team/pytorch_geometric/pull/10124))
+- Fixed `_recursive_config()` for `torch.nn.ModuleList` and `torch.nn.ModuleDict` ([#10124](https://github.com/pyg-team/pytorch_geometric/pull/10124), [#10129](https://github.com/pyg-team/pytorch_geometric/pull/10129))
 - Fixed the `k_hop_subgraph()` method for directed graphs ([#9756](https://github.com/pyg-team/pytorch_geometric/pull/9756))
 - Fixed `utils.group_cat` concatenating dimension ([#9766](https://github.com/pyg-team/pytorch_geometric/pull/9766))
 - Fixed `WebQSDataset.process` raising exceptions ([#9665](https://github.com/pyg-team/pytorch_geometric/pull/9665))

test/test_config_mixin.py

@@ -41,11 +41,18 @@ def __init__(self, p: float):
 
 @register(with_target=True)
 class CompoundModule(torch.nn.Module, ConfigMixin):
-    def __init__(self, z: int, module: Module, submodules: list[SubModule]):
+    def __init__(
+        self,
+        z: str,
+        module: Module,
+        submodules: list[SubModule],
+        key_modules: dict[str, torch.nn.Module],
+    ):
         super().__init__()
         self.z = z
         self.module = module
         self.submodules = torch.nn.ModuleList(submodules)
+        self.key_modules = torch.nn.ModuleDict(key_modules)
 
 
 def test_config_mixin() -> None:
@@ -107,28 +114,45 @@ def test_config_mixin() -> None:
 
 
 def test_config_mixin_compound() -> None:
-    submodules = [SubModule(1.41), SubModule(3.14)]
     module = Module(x=0, data=Dataclass(x=1, y=2))
-    model = CompoundModule(z=3, module=module, submodules=submodules)
+    submodules = [SubModule(1.41), SubModule(3.14)]
+    key_modules = {
+        "key1": Module(x=10, data=Dataclass(x=11, y=12)),
+        "key2": SubModule(2.71),
+    }
+    model = CompoundModule(z="foo", module=module, submodules=submodules,
+                           key_modules=key_modules)
+
     cfg = model.config()
     assert is_dataclass(cfg)
     assert cfg._target_ == 'test_config_mixin.CompoundModule'
-    assert cfg.z == 3
+    assert cfg.z == "foo"
     assert cfg.module._target_ == 'test_config_mixin.Module'
     assert cfg.module.x == 0
     assert isinstance(cfg.module.data, Dataclass)
     assert cfg.module.data.x == 1
     assert cfg.module.data.y == 2
+
     assert len(cfg.submodules) == 2
     assert isinstance(cfg.submodules, Sequence)
     assert cfg.submodules[0]._target_ == 'test_config_mixin.SubModule'
     assert cfg.submodules[0].p == 1.41
     assert cfg.submodules[1]._target_ == 'test_config_mixin.SubModule'
     assert cfg.submodules[1].p == 3.14
 
+    assert len(cfg.key_modules) == 2
+    assert cfg.key_modules["key1"]._target_ == 'test_config_mixin.Module'
+    assert cfg.key_modules["key1"].x == 10
+    assert isinstance(cfg.key_modules["key1"].data, Dataclass)
+    assert cfg.key_modules["key1"].data.x == 11
+    assert cfg.key_modules["key1"].data.y == 12
+
+    assert cfg.key_modules["key2"]._target_ == 'test_config_mixin.SubModule'
+    assert cfg.key_modules["key2"].p == 2.71
+
     model = CompoundModule.from_config(cfg)
     assert isinstance(model, CompoundModule)
-    assert model.z == 3
+    assert model.z == "foo"
     assert isinstance(model.module, Module)
     assert model.module.x == 0
     assert isinstance(model.module.data, Dataclass)
@@ -140,3 +164,12 @@ def test_config_mixin_compound() -> None:
     assert model.submodules[0].p == 1.41
     assert isinstance(model.submodules[1], SubModule)
     assert model.submodules[1].p == 3.14
+    assert isinstance(model.key_modules, torch.nn.ModuleDict)
+    assert len(model.key_modules) == 2
+    assert isinstance(model.key_modules["key1"], Module)
+    assert model.key_modules["key1"].x == 10
+    assert isinstance(model.key_modules["key1"].data, Dataclass)
+    assert model.key_modules["key1"].data.x == 11
+    assert model.key_modules["key1"].data.y == 12
+    assert isinstance(model.key_modules["key2"], SubModule)
+    assert model.key_modules["key2"].p == 2.71

torch_geometric/config_mixin.py

@@ -1,7 +1,9 @@
 import inspect
 from dataclasses import fields, is_dataclass
 from importlib import import_module
-from typing import Any, Dict, Iterable
+from typing import Any, Dict
+
+from torch.nn import ModuleDict, ModuleList
 
 from torch_geometric.config_store import (
     class_from_dataclass,
@@ -71,9 +73,9 @@ def _recursive_config(value: Any) -> Any:
         return value.config()
     if is_torch_instance(value, ConfigMixin):
         return value.config()
-    if isinstance(value, (tuple, list, Iterable)):
+    if isinstance(value, (tuple, list, ModuleList)):
         return [_recursive_config(v) for v in value]
-    if isinstance(value, dict):
+    if isinstance(value, (dict, ModuleDict)):
         return {k: _recursive_config(v) for k, v in value.items()}
     return value