• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1from __future__ import annotations
2
3from typing import Generic, TypeVar
4
5import torch
6
7__all__ = ['Await']
8
9W = TypeVar("W")
10
11class _PyAwaitMeta(type(torch._C._Await), type(Generic)):  # type: ignore[misc, no-redef]
12    pass
13
14class _Await(torch._C._Await, Generic[W], metaclass=_PyAwaitMeta):
15    r"""
16    Wrapper around a ``torch._C.Await`` which encapsulates delayed execution
17    of a callable. All manipulations happen with functions ``torch.jit._awaitable``,
18    ``torch.jit._awaitable_wait``, ``torch.jit._awaitable_nowait``.
19
20    Torch scriptable manipulations:
21    ``torch.jit._awaitable(func, *args)``
22    Creates ``Await[W]`` object, where W is return type of func.
23
24    Returns:
25    ``torch.jit._awaitable_wait(Await[W])``
26    Returns the result of the function, specified at ``_awaitable``,  with specified arguments.
27
28    Returns:
29        The result of type ``W`` of the function call. The result is owned by ``Await[W]``
30        and returned on all following ``_awaitable_wait`` calls.
31
32
33    ``torch.jit._awaitable_nowait(W)``
34    Returns:
35        Trivial ``Await[W]`` with specified result.
36
37
38    Only in eager mode:
39    ``fn() -> Callable[Tuple[Any], W]``
40    Returns:
41        Specified at ``_awaitable`` python function ``func``.
42
43    ``args() -> Tuple[Any]``
44    Returns:
45        Specified at ``_awaitable`` python args.
46
47    ``is_nowait() -> _bool``
48    Returns:
49        ``True`` if this object was created via ``_awaitable_nowait`` call (trivial `Await[W]`).
50
51    In eager mode ``Await[W]`` can be used as ``W`` i.e. attributes of W can be called on ``Await[W]``,
52    ``_awaitable_wait()`` call will be transparently added.
53    """
54