def forward(self, a: Tensor, dim: int, start: int, length: int):
return a.narrow(dim, start, length).clone()
)JIT";
+
+const auto two_tuple_unpack_script = R"JIT(
+ def forward(self, tup: Tuple[Tensor, Tensor]):
+ a, b = tup
+ return (a, b)
+)JIT";
+
+const auto three_tuple_unpack_script = R"JIT(
+ def forward(self, tup: Tuple[Tensor, Tensor, Tensor]):
+ a, b, c = tup
+ return (a, b, c)
+)JIT";
auto dtype = at::ScalarType::Int;
auto cpu = at::Device(DeviceType::CPU);
c10::List<int64_t> size0{4, 5};
- std::vector<IValue> args{
- size0, 4, dtype, at::kStrided, cpu, false};
+ std::vector<IValue> args{size0, 4, dtype, at::kStrided, cpu, false};
c10::List<int64_t> size1{5, 6};
- std::vector<IValue> args2{
- size1, 5, dtype, at::kStrided, cpu, false};
+ std::vector<IValue> args2{size1, 5, dtype, at::kStrided, cpu, false};
testStaticRuntime(full_script, args);
testStaticRuntime(full_script, args, args2);
}
testStaticRuntime(narrow_with_int_script, args_a);
testStaticRuntime(narrow_with_int_script, args_a, args_b);
}
+
+TEST(StaticRuntime, InvidualOps_TupleUnpack) {
+ auto two_tup = c10::ivalue::Tuple::create({at::randn({1}), at::randn({1})});
+ auto two_tup_large =
+ c10::ivalue::Tuple::create({at::randn({2, 2}), at::randn({2, 2})});
+
+ auto three_tup = c10::ivalue::Tuple::create(
+ {at::randn({1}), at::randn({1}), at::randn({1})});
+ auto three_tup_large = c10::ivalue::Tuple::create(
+ {at::randn({2, 2}), at::randn({2, 2}), at::randn({2, 2})});
+
+ testStaticRuntime(two_tuple_unpack_script, {two_tup});
+ testStaticRuntime(two_tuple_unpack_script, {two_tup}, {two_tup_large});
+
+ testStaticRuntime(three_tuple_unpack_script, {three_tup});
+ testStaticRuntime(three_tuple_unpack_script, {three_tup}, {three_tup_large});
+}
});
REGISTER_NATIVE_OPERATOR_FUNCTOR(
+ prim::TupleUnpack,
+ prim_TupleUnpack,
+ [](Node* n) -> SROperator {
+ return [](ProcessedNode* p_node) {
+ const auto& elems = p_node->Input(0).toTuple()->elements();
+ const size_t num_outputs = p_node->outputs().size();
+ TORCH_CHECK(
+ num_outputs == elems.size(),
+ "Number of outputs must match number of tuple elements.")
+ for (size_t i = 0; i < num_outputs; ++i) {
+ p_node->Output(i) = elems[i];
+ }
+ };
+ });
+
+REGISTER_NATIVE_OPERATOR_FUNCTOR(
prim::DictConstruct,
prim_DictConstruct,
[](Node* n) -> SROperator {