pytorch
5f1ac188 - add numpy typing plugin to mypy config (#92930)

Commit
1 year ago
add numpy typing plugin to mypy config (#92930) 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
Author
Committer
Parents
Loading