mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
3bf922a6ce
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/106249 Approved by: https://github.com/Skylion007
35 lines
866 B
Python
35 lines
866 B
Python
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"""
|
|
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")
|