Fix torch.nonzero type annotation (#51635)

rgommers · facebook-github-bot · commit 649e683255c1 · 2021-02-08T08:45:44.000-08:00
Summary: The overloads are a little tricky here. It's important that the overloads are such that it's unambiguous what `torch.nonzero(x)` will resolve to - so just specify defaults for one of the overloads. Also, `out` is left out of the second overload because a non-None value for `out` is not valid in combination with `as_tuple=True`. Closes gh-51434 Pull Request resolved: #51635 Reviewed By: zhangguanheng66 Differential Revision: D26279203 Pulled By: walterddr fbshipit-source-id: 8459c04fc9fbf7fc5f31b3f631aaac2f98b17ea6
diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py
@@ -322,8 +322,8 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) -
                  ' layout: _layout=strided, {}) -> Tensor: ...'
                  .format(FACTORY_PARAMS)],
         'is_grad_enabled': ['def is_grad_enabled() -> _bool: ...'],
-        'nonzero': ['def nonzero(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...',
-                    'def nonzero(input: Tensor, *, as_tuple: bool=...) -> Tensor: ...'],
+        'nonzero': ['def nonzero(input: Tensor, *, as_tuple: Literal[False]=False, out: Optional[Tensor]=None) -> Tensor: ...',
+                    'def nonzero(input: Tensor, *, as_tuple: Literal[True]) -> Tuple[Tensor, ...]: ...'],
         'binary_cross_entropy_with_logits': ['def binary_cross_entropy_with_logits(input: Tensor, target: Tensor, '
                                              'weight: Optional[Tensor] = None, size_average: Optional[bool] = None, '
                                              'reduce: Optional[bool] = None, reduction: str = ..., '
@@ -424,7 +424,8 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) -
         'element_size': ['def element_size(self) -> _int: ...'],
         'data_ptr': ['def data_ptr(self) -> _int: ...'],
         'dim': ['def dim(self) -> _int: ...'],
-        'nonzero': ['def nonzero(self, *, as_tuple: _bool=...) -> Tensor: ...'],
+        'nonzero': ['def nonzero(self, *, as_tuple: Literal[False]=False) -> Tensor: ...',
+                    'def nonzero(self, *, as_tuple: Literal[True]) -> Tuple[Tensor, ...]: ...'],
         'numel': ['def numel(self) -> _int: ...'],
         'ndimension': ['def ndimension(self) -> _int: ...'],
         'nelement': ['def nelement(self) -> _int: ...'],
diff --git a/torch/_C/_VariableFunctions.pyi.in b/torch/_C/_VariableFunctions.pyi.in
@@ -2,6 +2,7 @@
 
 from torch import Tensor, Generator, strided, memory_format, contiguous_format, strided
 from typing import List, Tuple, Optional, Union, Any, ContextManager, Callable, overload, Iterator, NamedTuple, Sequence, TypeVar
+from typing_extensions import Literal
 from torch._six import inf
 
 from torch.types import _int, _float, _bool, Number, _dtype, _device, _qscheme, _size, _layout
diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in
@@ -8,6 +8,7 @@ from typing import (
     Any, BinaryIO, Callable, ContextManager, Dict, Iterable, Iterator, List,
     NamedTuple, Optional, overload, Sequence, Tuple, TypeVar, Type, Union,
     Generic, Set, AnyStr)
+from typing_extensions import Literal
 from torch._six import inf
 
 from torch.types import _int, _float, _bool, _dtype, _device, _qscheme, _size, _layout, Device, Number, Storage
