1# mypy: allow-untyped-defs 2import torch 3from torch._jit_internal import _Await 4from torch.jit._builtins import _register_builtin 5from torch.utils import set_module 6 7 8set_module(_Await, "torch.jit") 9 10 11def _awaitable(func, *args, **kwargs): 12 r"""Create Await object that will call specified functioni with specified args, when it is requested for the result.""" 13 return torch._C._awaitable(func, *args, **kwargs) 14 15 16def _awaitable_wait(aw): 17 r"""Request await the result of execution, if Await is not completed yet, the func will be called immediately.""" 18 return torch._C._awaitable_wait(aw) 19 20 21def _awaitable_nowait(o): 22 r"""Create completed Await with specified result.""" 23 return torch._C._awaitable_nowait(o) 24 25 26_register_builtin(_awaitable_wait, "prim::awaitable_wait") 27_register_builtin(_awaitable_nowait, "prim::awaitable_nowait") 28