Is there any way to give a type hint to the output of torch.jit._wait?
The below example fails to compile to torchscript
import torch
from typing import List
def process(i: int) -> int:
return i + 1
@torch.jit.script
def process_many(l: List[int]) -> List[int]:
futs: List[torch.jit.Future] = []
out: List[int] = []
for v in l:
futs.append(torch.jit._fork(process, v))
for f in futs:
out.append(torch.jit._wait(f))
return out
with the error
File "/usr/local/lib/python3.7/site-packages/torch/jit/__init__.py", line 1281, in script
fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj))
RuntimeError:
Unknown type name 'torch.jit.Future':
File "./demo.py", line 10
@torch.jit.script
def process_many(l: List[int]) -> List[int]:
futs: List[torch.jit.Future] = []
~~~~~~~~~~~~~~~~ <--- HERE
out: List[int] = []
However, if i change the for loop to a list comprehension, it compiles fine
@torch.jit.script
def process_many(l: List[int]) -> List[int]:
futs = [torch.jit._fork(process, v) for v in l]
out: List[int] = []
for f in futs:
out.append(torch.jit._wait(f))
return out
# >>> print(process_many([0, 3, 5])
# [1, 4, 6]
What haven’t been able to figure out is if torch.jit.Future is the correct type annotation for me to be using here. For now, it seems possible to work around this just using a list comprehension, but it could get ugly if the logic becomes more complex.
I’m on 1.4.0, if it makes any difference