| # mypy: allow-untyped-defs | |
| import torch | |
| from torch._jit_internal import _Await | |
| from torch.jit._builtins import _register_builtin | |
| from torch.utils import set_module | |
| set_module(_Await, "torch.jit") | |
| def _awaitable(func, *args, **kwargs): | |
| r"""Create Await object that will call specified functioni with specified args, when it is requested for the result.""" | |
| return torch._C._awaitable(func, *args, **kwargs) | |
| def _awaitable_wait(aw): | |
| r"""Request await the result of execution, if Await is not completed yet, the func will be called immediately.""" | |
| return torch._C._awaitable_wait(aw) | |
| def _awaitable_nowait(o): | |
| r"""Create completed Await with specified result.""" | |
| return torch._C._awaitable_nowait(o) | |
| _register_builtin(_awaitable_wait, "prim::awaitable_wait") | |
| _register_builtin(_awaitable_nowait, "prim::awaitable_nowait") | |