mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
f3fce597e9
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter. You can review these PRs via: ```bash git diff --ignore-all-space --ignore-blank-lines HEAD~1 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/129769 Approved by: https://github.com/ezyang
28 lines
852 B
Python
28 lines
852 B
Python
# 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")
|