File size: 851 Bytes
c61ccee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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")