__init__.py 1.6 KB

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