mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
1b25fdbcd0
Summary: Added stubs for: * The `device` module * The `cuda` module * Parts of the `optim` module * Began adding stubs for the `autograd` module. I'll annotate more later but `no_grad` and friends are probably the most used exports from it so it seemed like a good place to start. This would close #16996, although comments on that issue reference other missing stubs so maybe it's worth keeping open as an umbrella issue. The big remaining missing package is `nn`. Also added a `py.typed` file so mypy will pick up on the type stubs. That closes #17639. Pull Request resolved: https://github.com/pytorch/pytorch/pull/18511 Differential Revision: D14715053 Pulled By: ezyang fbshipit-source-id: 9e4882ac997063650e6ce47604b3eaf1232c61c9
22 lines
707 B
Python
22 lines
707 B
Python
from typing import Any, Callable, TypeVar
|
|
|
|
# Used for annotating the decorator usage of 'no_grad' and 'enable_grad'.
|
|
# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators
|
|
FuncType = Callable[..., Any]
|
|
T = TypeVar('T', bound=FuncType)
|
|
|
|
class no_grad:
|
|
def __enter__(self) -> None: ...
|
|
def __exit__(self, *args: Any) -> bool: ...
|
|
def __call__(self, func: T) -> T: ...
|
|
|
|
class enable_grad:
|
|
def __enter__(self) -> None: ...
|
|
def __exit__(self, *args: Any) -> bool: ...
|
|
def __call__(self, func: T) -> T: ...
|
|
|
|
class set_grad_enabled:
|
|
def __init__(self, mode: bool) -> None: ...
|
|
def __enter__(self) -> None: ...
|
|
def __exit__(self, *args: Any) -> bool: ...
|