Repository: pyg-team/pytorch_geometric
Pull request #7294 (2023-05-04 11:56:51)

Allow tuples as keys in `ModuleDict`/`ParameterDict`

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

> :exclamation: Current head 796a053 differs from pull request most recent head 6067daa. Consider uploading reports for the commit 6067daa to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7294   +/-   ##
=======================================
  Coverage   91.64%   91.64%           
=======================================
  Files         438      438           
  Lines       24278    24296   +18     
=======================================
+ Hits        22249    22267   +18     
  Misses       2029     2029           
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7294?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/module\_dict.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7294?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZHVsZV9kaWN0LnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/parameter\_dict.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7294?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3BhcmFtZXRlcl9kaWN0LnB5) | `100.00% <100.00%> (ø)` | |

: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 e5f57ed66..73052b80b 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 support for tuples as keys in `ModuleDict`/`ParameterDict` ([#7294](https://github.com/pyg-team/pytorch_geometric/pull/7294))
 - Added `NodePropertySplit` transform for creating node-level splits using structural node properties ([#6894](https://github.com/pyg-team/pytorch_geometric/pull/6894))
 - Added an option to preserve directed graphs in `CitationFull` datasets ([#7275](https://github.com/pyg-team/pytorch_geometric/pull/7275))
 - Added support for `torch.sparse.Tensor` in `DataLoader` ([#7252](https://github.com/pyg-team/pytorch_geometric/pull/7252))
diff --git a/torch_geometric/nn/module_dict.py b/torch_geometric/nn/module_dict.py
index 09c3cab8c..122e03a30 100644
--- a/torch_geometric/nn/module_dict.py
+++ b/torch_geometric/nn/module_dict.py
@@ -1,15 +1,20 @@
-from typing import Iterable, Mapping, Optional, Tuple
+from typing import Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Module
 
+Key = Union[str, Tuple[str, ...]]
+
 
 # `torch.nn.ModuleDict` doesn't allow `.` to be used in key names.
 # This `ModuleDict` will support it by converting the `.` to `#` in the
 # internal representation and converts it back to `.` in the external
-# representation.
+# representation. It also allows passing tuples as keys.
 class ModuleDict(torch.nn.ModuleDict):
-    def __init__(self, modules: Optional[Mapping[str, Module]] = None):
+    def __init__(
+        self,
+        modules: Optional[Mapping[Union[str, Tuple[str, ...]], Module]] = None,
+    ):
         if modules is not None:  # Replace the keys in modules:
             modules = {
                 self.to_internal_key(key): module
@@ -18,28 +23,34 @@ class ModuleDict(torch.nn.ModuleDict):
         super().__init__(modules)
 
     @staticmethod
-    def to_internal_key(key: str) -> str:
-        return key.replace(".", "#")
+    def to_internal_key(key: Key) -> str:
+        if isinstance(key, tuple):
+            assert len(key) > 1
+            key = f"<{'___'.join(key)}>"
+        assert isinstance(key, str)
+        return key.replace('.', '#')
 
     @staticmethod
-    def to_external_key(key: str) -> str:
-        return key.replace("#", ".")
+    def to_external_key(key: str) -> Key:
+        key = key.replace('#', '.')
+        if key.startswith('<') and key.endswith('>') and '___' in key:
+            key = tuple(key[1:-1].split('___'))
+        return key
 
-    def __getitem__(self, key: str) -> Module:
+    def __getitem__(self, key: Key) -> Module:
         return super().__getitem__(self.to_internal_key(key))
 
-    def __setitem__(self, key: str, module: Module):
+    def __setitem__(self, key: Key, module: Module):
         return super().__setitem__(self.to_internal_key(key), module)
 
-    def __delitem__(self, key: str):
+    def __delitem__(self, key: Key):
         return super().__delitem__(self.to_internal_key(key))
 
-    def __contains__(self, key: str) -> bool:
+    def __contains__(self, key: Key) -> bool:
         return super().__contains__(self.to_internal_key(key))
 
-    def keys(self) -> Iterable[str]:
+    def keys(self) -> Iterable[Key]:
         return [self.to_external_key(key) for key in super().keys()]
 
-    def items(self) -> Iterable[Tuple[str, Module]]:
-        return [(self.to_external_key(key), value)
-                for key, value in super().items()]
+    def items(self) -> Iterable[Tuple[Key, Module]]:
+        return [(self.to_external_key(k), v) for k, v in super().items()]
diff --git a/torch_geometric/nn/parameter_dict.py b/torch_geometric/nn/parameter_dict.py
index 1d8ebc59b..99d887bd1 100644
--- a/torch_geometric/nn/parameter_dict.py
+++ b/torch_geometric/nn/parameter_dict.py
@@ -1,42 +1,57 @@
-from typing import Iterable, Mapping, Optional
+from typing import Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Parameter
 
+Key = Union[str, Tuple[str, ...]]
+
 
 # `torch.nn.ParameterDict` doesn't allow `.` to be used in key names.
 # This `ParameterDict` will support it by converting the `.` to `#` in the
 # internal representation and converts it back to `.` in the external
-# representation.
+# representation. It also allows passing tuples as keys.
 class ParameterDict(torch.nn.ParameterDict):
-    def __init__(self, parameters: Optional[Mapping[str, Parameter]] = None):
+    def __init__(
+        self,
+        parameters: Optional[Mapping[Key, Parameter]] = None,
+    ):
         # Replace the keys in modules.
         if parameters:
             parameters = {
-                self.to_internal_key(key): module
-                for key, module in parameters.items()
+                self.to_internal_key(key): parameter
+                for key, parameter in parameters.items()
             }
         super().__init__(parameters)
 
     @staticmethod
-    def to_internal_key(key: str) -> str:
-        return key.replace(".", "#")
+    def to_internal_key(key: str) -> Key:
+        if isinstance(key, tuple):
+            assert len(key) > 1
+            key = f"<{'___'.join(key)}>"
+        assert isinstance(key, str)
+        return key.replace('.', '#')
 
     @staticmethod
-    def to_external_key(key: str) -> str:
-        return key.replace("#", ".")
+    def to_external_key(key: str) -> Key:
+        key = key.replace('#', '.')
+        if key.startswith('<') and key.endswith('>') and '___' in key:
+            key = tuple(key[1:-1].split('___'))
+        return key
 
-    def __getitem__(self, key: str) -> Parameter:
+    def __getitem__(self, key: Key) -> Parameter:
         return super().__getitem__(self.to_internal_key(key))
 
-    def __setitem__(self, key: str, parameter: Parameter):
+    def __setitem__(self, key: Key, parameter: Parameter):
         return super().__setitem__(self.to_internal_key(key), parameter)
 
-    def __delitem__(self, key: str):
+    def __delitem__(self, key: Key):
         return super().__delitem__(self.to_internal_key(key))
 
-    def __contains__(self, key: str) -> bool:
+    def __contains__(self, key: Key) -> bool:
         return super().__contains__(self.to_internal_key(key))
 
-    def keys(self) -> Iterable[str]:
+    def keys(self) -> Iterable[Key]:
         return [self.to_external_key(key) for key in super().keys()]
+
+    def items(self) -> Iterable[Tuple[Key, Parameter]]:
+        return [(self.to_external_key(k), v) for k, v in super().items()]


Test patch:
diff --git a/test/nn/test_module_dict.py b/test/nn/test_module_dict.py
index 16fced0bc..6dbfc9fc6 100644
--- a/test/nn/test_module_dict.py
+++ b/test/nn/test_module_dict.py
@@ -7,9 +7,13 @@ def test_internal_external_key_conversion():
     assert ModuleDict.to_internal_key('a.b') == 'a#b'
     assert ModuleDict.to_internal_key('ab') == 'ab'
     assert ModuleDict.to_internal_key('a.b.c') == 'a#b#c'
+    assert ModuleDict.to_internal_key(('a', 'b')) == '<a___b>'
+    assert ModuleDict.to_internal_key(('a.b', 'c')) == '<a#b___c>'
 
     assert ModuleDict.to_external_key('a#b') == 'a.b'
     assert ModuleDict.to_external_key('a#b#c') == 'a.b.c'
+    assert ModuleDict.to_external_key('<a___b>') == ('a', 'b')
+    assert ModuleDict.to_external_key('<a#b___c>') == ('a.b', 'c')
 
 
 def test_dot_syntax_keys():
@@ -26,5 +30,24 @@ def test_dot_syntax_keys():
     for key in expected_keys:
         assert key in module_dict
 
+    assert 'model.lin2' in module_dict
     del module_dict['model.lin2']
     assert 'model.lin2' not in module_dict
+
+
+def test_tuple_keys():
+    module_dict = ModuleDict({
+        ('a', 'b'): torch.nn.Linear(16, 16),
+        ('a.b', 'c'): torch.nn.Linear(8, 8),
+    })
+
+    expected_keys = {('a', 'b'), ('a.b', 'c')}
+    assert set(module_dict.keys()) == expected_keys
+    assert set([key for key, _ in module_dict.items()]) == expected_keys
+
+    for key in expected_keys:
+        assert key in module_dict
+
+    assert ('a', 'b') in module_dict
+    del module_dict['a', 'b']
+    assert ('a', 'b') not in module_dict
diff --git a/test/nn/test_parameter_dict.py b/test/nn/test_parameter_dict.py
index 34cde8c54..29d44f926 100644
--- a/test/nn/test_parameter_dict.py
+++ b/test/nn/test_parameter_dict.py
@@ -1,33 +1,55 @@
-from typing import Mapping
-
 import torch
-from torch.nn import Parameter
 
 from torch_geometric.nn.parameter_dict import ParameterDict
 
 
 def test_internal_external_key_conversion():
-    assert ParameterDict.to_internal_key("a.b") == "a#b"
-    assert ParameterDict.to_internal_key("ab") == "ab"
-    assert ParameterDict.to_internal_key("a.b.c") == "a#b#c"
+    assert ParameterDict.to_internal_key('a.b') == 'a#b'
+    assert ParameterDict.to_internal_key('ab') == 'ab'
+    assert ParameterDict.to_internal_key('a.b.c') == 'a#b#c'
+    assert ParameterDict.to_internal_key(('a', 'b')) == '<a___b>'
+    assert ParameterDict.to_internal_key(('a.b', 'c')) == '<a#b___c>'
 
-    assert ParameterDict.to_external_key("a#b") == "a.b"
-    assert ParameterDict.to_external_key("a#b#c") == "a.b.c"
+    assert ParameterDict.to_external_key('a#b') == 'a.b'
+    assert ParameterDict.to_external_key('a#b#c') == 'a.b.c'
+    assert ParameterDict.to_external_key('<a___b>') == ('a', 'b')
+    assert ParameterDict.to_external_key('<a#b___c>') == ('a.b', 'c')
 
 
 def test_dot_syntax_keys():
-    parameters: Mapping[str, Parameter] = {
-        "param1": Parameter(torch.Tensor(16, 16)),
-        "model.param2": Parameter(torch.Tensor(8, 8)),
-        "model.sub_model.param3": Parameter(torch.Tensor(4, 4)),
+    parameter_dict = {
+        'param1': torch.nn.Parameter(torch.Tensor(16, 16)),
+        'model.param2': torch.nn.Parameter(torch.Tensor(8, 8)),
+        'model.sub_model.param3': torch.nn.Parameter(torch.Tensor(4, 4)),
+    }
+    parameter_dict = ParameterDict(parameter_dict)
+
+    expected_keys = {'param1', 'model.param2', 'model.sub_model.param3'}
+    assert set(parameter_dict.keys()) == expected_keys
+    assert set([key for key, _ in parameter_dict.items()]) == expected_keys
+
+    for key in expected_keys:
+        assert key in parameter_dict
+
+    assert 'model.param2' in parameter_dict
+    del parameter_dict['model.param2']
+    assert 'model.param2' not in parameter_dict
+
+
+def test_tuple_keys():
+    parameter_dict = {
+        ('a', 'b'): torch.nn.Parameter(torch.Tensor(16, 16)),
+        ('a.b', 'c'): torch.nn.Parameter(torch.Tensor(8, 8)),
     }
-    parameter_dict = ParameterDict(parameters)
+    parameter_dict = ParameterDict(parameter_dict)
 
-    expected_keys = {"param1", "model.param2", "model.sub_model.param3"}
+    expected_keys = {('a', 'b'), ('a.b', 'c')}
     assert set(parameter_dict.keys()) == expected_keys
+    assert set([key for key, _ in parameter_dict.items()]) == expected_keys
 
     for key in expected_keys:
         assert key in parameter_dict
 
-    del parameter_dict["model.param2"]
-    assert "model.param2" not in parameter_dict
+    assert ('a', 'b') in parameter_dict
+    del parameter_dict['a', 'b']
+    assert ('a', 'b') not in parameter_dict