Add __module__ to torch.dtype.__dict__ (#65182)
Summary:
torch.dtype.__reduce__ returns a string, which causes Pickle to look
up the object by module and name. In order to find the right module,
Pickle looks for __module__ on the object; if it doesn't find that, it
falls back to searching sys.modules.
Previously, torch.dtype instances did not have a `__module__`
attribute, so pickling dtypes would fall back to a search of
sys.module.
Instances of normal Python objects have a `__module__` attribute
because normal Python classes have a `__module__` key in their
`__dict__`. Imitate that by populating one in `torch.dtype`.
We set the field in `tp_dict` before calling `PyType_Ready` (instead
of afterwards) because of the doc warning against mutating a type's
dictionary once initialized:
https://docs.python.org/3/c-api/typeobj.html#c.PyTypeObject.tp_dict
fixes https://github.com/pytorch/pytorch/issues/65077
---
I didn't add any tests because I didn't see any obvious places with similar tests for pickling or dtype objects. Let me know if I missed the right place, or should start one.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/65182
Reviewed By: mrshenli
Differential Revision: D31310530
Pulled By: ezyang
fbshipit-source-id: 20cd713ce175a709d6ce47459c3891162ce29d77
