op

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -147,7 +147,7 @@ def test_op_cast(self):
             ir_version=9,
             opset_imports=[oh.make_opsetid("", 18)],
         )
-        self._finalize_test(model, (torch.rand((4, 5, 6, 7), dtype=torch.float32),))
+        self._finalize_test(model, torch.rand((4, 5, 6, 7), dtype=torch.float32))
 
     def test_op_transpose(self):
         model = oh.make_model(
@@ -180,6 +180,46 @@ def test_op_reshape(self):
             model, torch.rand((4, 5, 6, 7), dtype=torch.float32), torch.tensor([7, 4, 6, 5])
         )
 
+    def test_op_matmul(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("MatMul", ["X", "Y"], ["Z"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c", "d"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "d", "f"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c", "f"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand((4, 5, 6, 7), dtype=torch.float32),
+            torch.rand((4, 5, 7, 11), dtype=torch.float32),
+        )
+
+    def test_op_unsqueeze(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Unsqueeze", ["X", "axes"], ["Z"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", 1, "d"]),
+                    oh.make_tensor_value_info("axes", TINT64, ["s"]),
+                ],
+                [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.tensor([2], dtype=torch.int64),
+        )
+
 
 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, Mul_1, Sub_1
+from .binary_ops import Add_1, Div_1, MatMul_1, Mul_1, Sub_1
 from .other_ops import Cast_6, Transpose_1
-from .shape_ops import Reshape_5, Shape_15, Squeeze_13
+from .shape_ops import Reshape_5, Shape_15, Squeeze_13, Unsqueeze_13

onnx_diagnostic/reference/torch_ops/binary_ops.py

@@ -8,18 +8,25 @@ def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
         return OpRunValue(x.tensor + y.tensor)
 
 
-class Mul_1(OpRun):
-    """Mul"""
+class Div_1(OpRun):
+    """Div"""
 
     def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
-        return OpRunValue(x.tensor * y.tensor)
+        return OpRunValue(x.tensor / y.tensor)
 
 
-class Div_1(OpRun):
-    """Div"""
+class MatMul_1(OpRun):
+    """MatMul"""
 
     def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
-        return OpRunValue(x.tensor / y.tensor)
+        return OpRunValue(x.tensor @ y.tensor)
+
+
+class Mul_1(OpRun):
+    """Mul"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor * y.tensor)
 
 
 class Sub_1(OpRun):

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -8,7 +8,9 @@
 class Cast_6(OpRun):
     def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         super().__init__(node, version)
-        self.to = onnx_dtype_to_torch_dtype(self.get_attribute_int(node, "to", 0))
+        to = self.get_attribute_int(node, "to", 0)
+        assert isinstance(to, int), f"Unexpected value for attribute to={to!r}"
+        self.to = onnx_dtype_to_torch_dtype(to)
 
     def run(self, data: OpRunValue) -> OpRunValue:
         return OpRunValue(data.tensor.to(self.to))

onnx_diagnostic/reference/torch_ops/shape_ops.py

@@ -28,3 +28,13 @@ class Squeeze_13(OpRun):
 
     def run(self, data: OpRunValue, axes: OpRunValue) -> OpRunValue:
         return OpRunValue(data.tensor.squeeze(axes.as_tuple_int))
+
+
+class Unsqueeze_13(OpRun):
+    "Unsqueeze"
+
+    def run(self, data: OpRunValue, axes: OpRunValue) -> OpRunValue:
+        t = data.tensor
+        for i in axes.as_tuple_int:
+            t = t.unsqueeze(i)
+        return OpRunValue(t)