| import torch |
| |
| from torch.utils import set_module |
| from torch.jit._builtins import _register_builtin |
| from torch._jit_internal import _Await |
| |
| set_module(_Await, "torch.jit") |
| |
| def _awaitable(func, *args, **kwargs): |
| r""" |
| Creates 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""" |
| Requests 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""" |
| Creates 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") |
