Repository: pyg-team/pytorch_geometric
Pull request #10129 (2025-03-20 10:35:32)

Fix for configmixin with explicit classes
Fix for configmixin with explicit classes


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b76cf93ba..19646543e 100644
--- a/CHANGELOG.md
+++ b/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))
diff --git a/torch_geometric/config_mixin.py b/torch_geometric/config_mixin.py
index f6e357eb8..b83d8ee28 100644
--- a/torch_geometric/config_mixin.py
+++ b/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
 


Test patch:
diff --git a/test/test_config_mixin.py b/test/test_config_mixin.py
index ae146a23c..a9f31524f 100644
--- a/test/test_config_mixin.py
+++ b/test/test_config_mixin.py
@@ -41,11 +41,18 @@ class SubModule(Base):
 
 @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,18 +114,25 @@ 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'
@@ -126,9 +140,19 @@ def test_config_mixin_compound() -> None:
     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