y = torch.jit._wait(fut)
# assert nothing; only to make sure the fake python path works
+ def test_async_parsing(self):
+ @torch.jit.script
+ def foo(x):
+ # type: (Tensor) -> List[Tensor]
+ return [torch.neg(x), x.t()]
+
+ @torch.jit.script
+ def bar(x):
+ futures = torch.jit.annotate(List[Future[List[Tensor]]], [])
+ for _ in range(3):
+ future = torch.jit.annotate(
+ Future[List[Tensor]],
+ torch.jit._fork(foo, x)
+ )
+ futures.append(future)
+
+ output = torch.jit.annotate(List[List[Tensor]], [])
+ for i in range(3):
+ output.append(torch.jit._wait(futures[i]))
+ return output
+
+ x = torch.rand(3, 3)
+ result = bar(x)
+ self.assertEqual(len(result), 3)
+
def test_async_script(self):
@torch.jit.script
def foo(x):
auto elem_type = parseTypeFromExpr(*subscript.subscript_exprs().begin());
return OptionalType::create(elem_type);
}},
+ {"Future", [](Subscript subscript) -> TypePtr {
+ if (subscript.subscript_exprs().size() != 1) {
+ throw ErrorReport(subscript) << " expected exactly one element type but found " << subscript.subscript_exprs().size();
+ }
+ auto elem_type = parseTypeFromExpr(*subscript.subscript_exprs().begin());
+ return FutureType::create(elem_type);
+ }},
};
return map;
}