Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add numpy typing plugin to mypy config (pytorch#92930)
This added the numpy typing plugin to mypy config so that we could use it for DeviceMesh typing annotations Please see pytorch#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: pytorch#92930 Approved by: https://github.com/ezyang, https://github.com/malfet
- Loading branch information