Improves type-checking guards. (#43339)
Summary:
PR https://github.com/pytorch/pytorch/issues/38157 fixed type checking for mypy by including `if False` guards on some type-checker-only imports. However other typecheckers - [like pyright](https://github.com/microsoft/pylance-release/issues/262#issuecomment-677758245) - will respect this logic and ignore the imports. Using [`if TYPE_CHECKING`](https://docs.python.org/3/library/typing.html#typing.TYPE_CHECKING) instead means both mypy and pyright will work correctly.
[For background, an example of where the current code fails](https://github.com/microsoft/pylance-release/issues/262) is if you make a file `tmp.py` with the contents
```python
import torch
torch.ones((1,))
```
Then [`pyright tmp.py --lib`](https://github.com/microsoft/pyright#command-line) will fail with a `"ones" is not a known member of module` error. This is because it can't find the `_VariableFunctions.pyi` stub file, as pyright respects the `if False` logic. After adding the `TYPE_CHECKING` guard, all works correctly.
Credit to erictraut for suggesting the fix.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43339
Reviewed By: agolynski
Differential Revision: D23348142
Pulled By: ezyang
fbshipit-source-id: c8a58122a7b0016845c311da39a1cc48748ba03f
