[cpp_extension] Add abstract base class to OrderedDictWrapper, and clarify torch.nn.Module typing #160897

benjaminglass1 · 2025-08-18T17:30:15Z

Stack from ghstack (oldest at bottom):

-> [cpp_extension] Add abstract base class to OrderedDictWrapper, and clarify torch.nn.Module typing #160897

cc @malfet @zou3519 @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

[ghstack-poisoned]

pytorch-bot · 2025-08-18T17:30:19Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160897

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b8dc3d8 with merge base 1c16c18 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

torch/csrc/api/src/python/init.cpp

benjaminglass1 · 2025-08-18T17:31:47Z

torch/nn/modules/module.py

+        replica._buffers = dict(replica._buffers)
+        replica._modules = dict(replica._modules)


MutableMapping doesn't define .copy. This was the only code change that was required to make this change work.

[ghstack-poisoned]

…arify torch.nn.Module typing ghstack-source-id: bd07986 Pull Request resolved: #160897

eellison · 2025-09-03T18:31:14Z

cc @albanD mind taking a look?

torch/nn/cpp.py

[ghstack-poisoned]

…arify torch.nn.Module typing ghstack-source-id: fd1955f Pull Request resolved: #160897

albanD · 2025-09-11T20:05:54Z

torch/nn/modules/module.py

    training: bool
-    _parameters: dict[str, Optional[Parameter]]
-    _buffers: dict[str, Optional[Tensor]]
+    _parameters: MutableMapping[str, Optional[Parameter]]


Why change this? All of these are actually dict right?

@albanD Only in some cases! For instance, they could be OrderedDictWrapper, which wraps a pybind11 C++ dict-like class.

I don't think this is really used in practice but ok.
I'll let @lolpack decide on this one since it's pure typing.

It's hard to say what the downstream effect of this change is, but if the underlying value is not always a dict than this seems like a legitimate increase in strictness.

Sorry to kick the question back to you, but it's hard to say what the downstream impact on users would be without more context. If this is not user facing, then IMO you might as well switch. If it is user facing, then I think you need to trade off between backwards-compat, convenience, and safety.

Maybe you know that in all user-facing cases, the value is actually a dict, and users can safely use dict-only methods? In that case, I say continue using dict.

If the MutableMapping interface is sufficient for common use cases, maybe it's worth changing it for safety, even at the risk of some churn for users.

I think it's a judgment call.

rec · 2025-09-23T08:23:45Z

I ran a "delta" on the pyright type checking (using this utility) before and after this commit: the change is not huge, but universally positive, adding +0.12% to "completenessScore".

[
    {
        "absolute": {
            "exportedSymbolCounts": {
                "withKnownType": 8,
                "withUnknownType": -6
            },
            "completenessScore": 0.0004445407249853428,
        },
        "percent": {
            "exportedSymbolCounts": {
                "withKnownType": 0.13054830287206268,
                "withUnknownType": -0.06430178973314757
            },
            "completenessScore": 0.11826611356782056,
        },
        "symbols": {
            "added": [
                "torch.nn.cpp.OrderedDictWrapper.__delitem__",
                "torch.nn.cpp.OrderedDictWrapper.__setitem__"
            ],
            "common": {
                "torch.nn.cpp.OrderedDictWrapper.__contains__": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing",
                            "Type annotation for parameter \"key\" is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.__getitem__": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing",
                            "Type annotation for parameter \"key\" is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.__init__": {
                    "diagnostics": {
                        "removed": [
                            "Type annotation for parameter \"attr\" is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.__iter__": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.__len__": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.cpp_module": {
                    "isTypeKnown": false,
                    "diagnostics": {
                        "added": [
                            "Type unknown for variable \"torch.nn.cpp.OrderedDictWrapper.cpp_module\""
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.items": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.keys": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing"
                        ]
                    }
                },
                "torch.nn.cpp.OrderedDictWrapper.values": {
                    "isTypeKnown": true,
                    "diagnostics": {
                        "removed": [
                            "Return type annotation is missing"
                        ]
                    }
                }
            },
            "removed": []
        },
        "file_names": [
            "/home/rec/git/pytorch/outputs/pyright.1c16c18a534.json",
            "/home/rec/git/pytorch/outputs/pyright.11e97bc7bd4.json"
        ]
    }
]

Skylion007 · 2025-09-30T19:26:03Z

torch/csrc/api/src/python/init.cpp

 namespace torch::python {
 namespace {
 template <typename T>
 void bind_ordered_dict(py::module module, const char* dict_name) {


Ah, this is going to break cpp20 because module is a keyword, fun...

Skylion007 · 2025-09-30T19:27:01Z

torch/nn/cpp.py

    # must manually override them.

-    def items(self):
+    def items(self) -> list[tuple[str, _T]]:  # type: ignore[override]


Suggested change

def items(self) -> list[tuple[str, _T]]: # type: ignore[override]

def items(self) -> collections.abc.ItemsView[str, _T]]:

@Skylion007 So this would actually be incorrect. What is returned from this method is not an ItemsView, but a list of tuple. We'd have to rewrite parts of the underlying C++ class to make an ItemsView possible, is my read after tackling this code. Same applies to the two comments below.

Skylion007 · 2025-09-30T19:27:25Z

torch/nn/cpp.py

        return self.cpp_dict.items()

-    def keys(self):
+    def keys(self) -> list[str]:  # type: ignore[override]


Suggested change

def keys(self) -> list[str]: # type: ignore[override]

def keys(self) -> collections.abc.KeysView[str]:

Skylion007 · 2025-09-30T19:27:43Z

torch/nn/cpp.py

        return self.cpp_dict.keys()

-    def values(self):
+    def values(self) -> list[_T]:  # type: ignore[override]


Suggested change

def values(self) -> list[_T]: # type: ignore[override]

def values(self) -> collections.abc.ValuesView[_T]:

Skylion007 · 2025-09-30T19:28:29Z

torch/nn/cpp.py

-    def __iter__(self):
+    # This should return an Iterator[str], but OrderedDict::item is not currently
+    # designed to let us iterate over only the keys.
+    def __iter__(self) -> Iterator[tuple[str, _T]]:  # type: ignore[override]


Why can't we just call iter keys? or is it an API change...

No API change; this wraps an underlying C++ custom class we've implemented, and we don't have methods available for simply looking at the keys. One could possibly be implemented, but I was more interested in documenting the existing behavior in this PR.

Update

c9ec2d6

[ghstack-poisoned]

This was referenced Aug 18, 2025

Update pybind11 submodule to 3.0.1 #160754

Closed

Update pybind11 enums to new native_enum class #160771

Closed

pytorch-bot bot added the release notes: cpp release notes category label Aug 18, 2025

benjaminglass1 mentioned this pull request Aug 18, 2025

Improve uncached FakeTensor propagation speed by 10% in some cases #160772

Draft

benjaminglass1 self-assigned this Aug 18, 2025

benjaminglass1 commented Aug 18, 2025

View reviewed changes

pytorchbot added the open source label Aug 18, 2025

benjaminglass1 added 2 commits August 19, 2025 19:44

Update

a53d85a

[ghstack-poisoned]

Update

d936829

[ghstack-poisoned]

benjaminglass1 marked this pull request as ready for review August 19, 2025 21:44

benjaminglass1 requested review from albanD, jbschlosser and mikaylagawarecki as code owners August 19, 2025 21:44

benjaminglass1 added 3 commits August 19, 2025 23:09

Update

6c2d23c

[ghstack-poisoned]

Update

77126c9

[ghstack-poisoned]

Update

9578f85

[ghstack-poisoned]

benjaminglass1 mentioned this pull request Aug 20, 2025

[cpp_wrapper] Swap to new PyBind11 simple GIL header #161063

Closed

benjaminglass1 added a commit that referenced this pull request Aug 20, 2025

[cpp_extension] Add abstract base class to OrderedDictWrapper, and cl…

947f0e0

…arify torch.nn.Module typing ghstack-source-id: bd07986 Pull Request resolved: #160897

Skylion007 reviewed Sep 9, 2025

View reviewed changes

torch/nn/cpp.py Outdated Show resolved Hide resolved

amjames added module: cpp-extensions Related to torch.utils.cpp_extension module: nn Related to torch.nn labels Sep 9, 2025

Update

b8dc3d8

[ghstack-poisoned]

benjaminglass1 added a commit that referenced this pull request Sep 10, 2025

[cpp_extension] Add abstract base class to OrderedDictWrapper, and cl…

11e97bc

…arify torch.nn.Module typing ghstack-source-id: fd1955f Pull Request resolved: #160897

benjaminglass1 requested a review from Skylion007 September 10, 2025 17:28

albanD reviewed Sep 11, 2025

View reviewed changes

Skylion007 reviewed Sep 30, 2025

View reviewed changes

		replica._buffers = dict(replica._buffers)
		replica._modules = dict(replica._modules)

	def items(self) -> list[tuple[str, _T]]: # type: ignore[override]
	def items(self) -> collections.abc.ItemsView[str, _T]]:

	def keys(self) -> list[str]: # type: ignore[override]
	def keys(self) -> collections.abc.KeysView[str]:

	def values(self) -> list[_T]: # type: ignore[override]
	def values(self) -> collections.abc.ValuesView[_T]:

[cpp_extension] Add abstract base class to OrderedDictWrapper, and clarify torch.nn.Module typing #160897

Are you sure you want to change the base?

[cpp_extension] Add abstract base class to OrderedDictWrapper, and clarify torch.nn.Module typing #160897

Uh oh!

Conversation

benjaminglass1 commented Aug 18, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160897

✅ No Failures

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison commented Sep 3, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rec commented Sep 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

benjaminglass1 commented Aug 18, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Aug 18, 2025 •

edited

Loading