add concat

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -220,6 +220,26 @@ def test_op_unsqueeze(self):
             torch.tensor([2], dtype=torch.int64),
         )
 
+    def test_op_concat(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Concat", ["X", "Y"], ["Z"], axis=2)],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", 1, "d"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", 1, "d"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "d"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand((4, 5, 1, 7), dtype=torch.float32),
+            torch.rand((4, 5, 2, 7), dtype=torch.float32),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -1,5 +1,5 @@
 from ._op_run import OpRun, OpRunValue
 from .access_ops import Slice_13
 from .binary_ops import Add_1, Div_1, MatMul_1, Mul_1, Sub_1
-from .other_ops import Cast_6, Transpose_1
+from .other_ops import Cast_6, Concat_1, Transpose_1
 from .shape_ops import Reshape_5, Shape_15, Squeeze_13, Unsqueeze_13

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -16,6 +16,17 @@ def run(self, data: OpRunValue) -> OpRunValue:
         return OpRunValue(data.tensor.to(self.to))
 
 
+class Concat_1(OpRun):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        axis = self.get_attribute_int(node, "axis", 0)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+
+    def run(self, *data: OpRunValue) -> OpRunValue:
+        return OpRunValue(torch.cat([t.tensor for t in data], axis=self.axis))
+
+
 class Transpose_1(OpRun):
     def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         super().__init__(node, version)