[ONNX] Add support for operator add between tensor list (#41888)

BowenBao · facebook-github-bot · commit da70976e661a · 2020-08-20T22:38:23.000-07:00
Summary: E.g. ```python outs = [] outs += [torch.randn(3,4)] outs = outs + [torch.randn(4,5), torch.randn(5,6)] ``` Pull Request resolved: #41888 Reviewed By: houseroad Differential Revision: D23172880 Pulled By: bzinodev fbshipit-source-id: 93865106e3de5908a993e0cfa82f626ba94dab7e
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -2700,11 +2700,15 @@ def forward(self, x):
                 res1 = []
                 arr = x.split([3, 4, 1, 1, 2, 3, 2], 0)
                 res2 = torch.zeros(3, 4, dtype=torch.long)
+                res3 = []
+                res4 = []
                 for i in range(len(arr)):
                     res = res.append(arr[i].sum(0, False))
                     res1 = res1.append(arr[-1 - i].sum(0, False))
                     res2 += 1
-                return torch.stack(res), torch.stack(res1), res2
+                    res3 = res3 + [arr[i].sum(0, False)]
+                    res4 += [arr[-1 - i].sum(0, False)]
+                return torch.stack(res), torch.stack(res1), res2, torch.stack(res3), torch.stack(res4)
 
         model = ListLoopModel()
         inputs = torch.randn(16)
@@ -2723,6 +2727,8 @@ def forward(self, x):
 
                 res.insert(0, tensors[1])
                 res.append(tensors[2])
+                res += [tensors[3], tensors[4]]
+                res = res + [tensors[5]]
                 return torch.ones(len(res))
 
         model = ListModel()
diff --git a/torch/onnx/symbolic_opset11.py b/torch/onnx/symbolic_opset11.py
@@ -284,6 +284,19 @@ def append(g, self, tensor):
     return g.op("SequenceInsert", self, tensor)
 
 
+def add(g, self, other, alpha=None):
+    if sym_help._is_value(self) and sym_help._is_tensor_list(self):
+        tensor_list_node = other.node()
+        if tensor_list_node.kind() != "prim::ListConstruct":
+            return _unimplemented("add", "does not support adding dynamic tensor list to another")
+        tensors = sym_help._unpack_list(other)
+        l = self
+        for t in tensors:
+            l = g.op("SequenceInsert", l, t)
+        return l
+
+    return torch.onnx.symbolic_opset9.add(g, self, other, alpha)
+
 def insert(g, self, pos, tensor):
     return g.op("SequenceInsert", self, tensor, pos)
 
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -83,6 +83,9 @@ def reshape_as(g, self, other):
 
 
 def add(g, self, other, alpha=None):
+    if sym_help._is_value(self) and sym_help._is_tensor_list(self):
+        return sym_help._onnx_opset_unsupported_detailed('Add', 9, 11, 'Add between list of tensors not supported')
+
     # default alpha arg is to allow no-alpha add (aten add st overload no alpha)
     if alpha and sym_help._scalar(sym_help._maybe_get_scalar(alpha)) != 1:
         return _unimplemented("add", "alpha != 1")
