Improve Tensor type hints

tools/pyi/gen_pyi.py

@@ -275,7 +275,7 @@ def generate_type_hints(fname, decls, is_tensor=False):
                 python_args.append('*')
                 render_kw_only_separator = False
             python_args += ["dtype: _dtype=None",
-                            "layout: layout=strided",
+                            "layout: _layout=strided",
                             "device: Union[_device, str, None]=None",
                             "requires_grad:_bool=False"]
 

torch/__init__.pyi.in

@@ -1,6 +1,6 @@
 # ${generated_comment}
 
-from typing import List, Tuple, Optional, Union, Any, ContextManager, Callable, overload
+from typing import List, Tuple, Optional, Union, Any, ContextManager, Callable, overload, Iterator
 from torch._six import inf
 
 import builtins
@@ -66,6 +66,7 @@ _dtype = dtype
 _device = device
 _qscheme = qscheme
 _size = Union[Size, List[_int], Tuple[_int, ...]]
+_layout = layout
 
 # Meta-type for "numeric" things; matches our docs
 Number = Union[builtins.int, builtins.float, builtins.bool]
@@ -83,16 +84,31 @@ class Generator:
 # torch.tensor.Tensor doesn't get type annotations.  Nobody
 # should really do that, so maybe this is not so bad.
 class Tensor:
-    dtype: _dtype = ...
-    shape: Size = ...
-    device: _device = ...
     requires_grad: _bool = ...
     grad: Optional[Tensor] = ...
     data: Tensor = ...
+    names: List[str] = ...
+    @property
+    def dtype(self) -> _dtype: ...
+    @property
+    def shape(self) -> Size: ...
+    @property
+    def device(self) -> _device: ...
+    @property
+    def T(self) -> Tensor: ...
+    @property
+    def grad_fn(self) -> Optional[Any]: ...
+    @property
+    def ndim(self) -> _int: ...
+    @property
+    def layout(self) -> _layout: ...
 
     ${tensor_method_hints}
 
     # Manually defined methods from torch/tensor.py
+    def __len__(self) -> _int: ...
+    def __iter__(self) -> Iterator[Tensor]: ...
+    def __contains__(self, item: Union[Tensor, Number]) -> _bool: ...
     def register_hook(self, hook: Callable) -> Any: ...
     def retain_grad(self) -> None: ...
     def is_shared(self) -> _bool: ...