frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

add numpy typing plugin to mypy config (#92930)

Browse Source 
This added the numpy typing plugin to mypy config so that we could
use it for DeviceMesh typing annotations

Please see https://github.com/pytorch/pytorch/pull/92931 about why we need this. For example, we are currently saving the DeviceMesh's mesh field as torch.Tensor, where when we do sth like:
```python
with FakeTensorMode():
    device_mesh = DeviceMesh("cuda", torch.arange(4))
```
It would throw error because FakeTensorMode or any TorchDispatchMode tracks every tensor creation and interactions. While DeviceMesh just want to save a nd-array to record the mesh topology, and would like to avoid the interaction with subsystems like FakeTensor, so we want to support saving `mesh` as numpy array instead.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/92930
Approved by: https://github.com/ezyang, https://github.com/malfet
This commit is contained in:
Wanchao Liang
2023-01-30 16:27:07 +00:00
committed by PyTorch MergeBot
parent 2a6e085704
commit 5f1ac188f8
3 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mypy-strict.ini
View File

@ -7,7 +7,7 @@
[mypy]
python_version = 3.8
plugins = mypy_plugins/check_mypy_version.py
plugins = mypy_plugins/check_mypy_version.py, numpy.typing.mypy_plugin
cache_dir = .mypy_cache/strict
strict_optional = True