fix a few things

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -435,6 +435,23 @@ def test_op_reduce_min(self):
             atol=1e-6,
         )
 
+    def test_op_reduce_min_no_axes(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceMin", ["X"], ["Z"], keepdims=0)],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            atol=1e-6,
+        )
+
     def test_op_reduce_min_17(self):
         model = oh.make_model(
             oh.make_graph(
@@ -452,6 +469,23 @@ def test_op_reduce_min_17(self):
             atol=1e-6,
         )
 
+    def test_op_reduce_min_17_no_axes(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceMin", ["X"], ["Z"])],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 17)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            atol=1e-6,
+        )
+
     def test_op_reduce_sum(self):
         model = oh.make_model(
             oh.make_graph(
@@ -1123,6 +1157,42 @@ def test_conv(self):
         B = torch.tensor([[[[0]]]], dtype=torch.float32)
         self._finalize_test(model, X, W, B, use_ort=True)
 
+    def test_conv_autopad_valid(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "Conv",
+                        ["X", "W", "B"],
+                        ["Y"],
+                        dilations=[1, 1],
+                        strides=[2, 2],
+                        auto_pad="VALID",
+                    )
+                ],
+                "g",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [None, None, None, None]),
+                    oh.make_tensor_value_info("W", TFLOAT, [None, None, None, None]),
+                    oh.make_tensor_value_info("B", TFLOAT, [None, None, None, None]),
+                ],
+                [oh.make_tensor_value_info("Y", TFLOAT, [None, None, None, None])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=10,
+        )
+        sH, sW = 5, 5
+        i = sH // 2
+        j = sW // 2
+        X = torch.zeros((1, 1, sH, sW), dtype=torch.float32)
+        X[0, 0, i, j] = 1.0
+        W = torch.zeros((1, 1, 3, 3), dtype=torch.float32)
+        W[0, 0, :, :] = torch.minimum(
+            2 ** torch.arange(9).reshape((3, -1)), torch.tensor([256])
+        )
+        B = torch.tensor([[[[0]]]], dtype=torch.float32)
+        self._finalize_test(model, X, W, B, use_ort=True)
+
     def test_nonzero(self):
         model = oh.make_model(
             oh.make_graph(

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -28,8 +28,11 @@ class OpRunTensor(OpRunValue):
     """
 
     def __init__(self, tensor, is_constant: bool = False, may_cpu: bool = False):
-        assert isinstance(tensor, torch.Tensor), f"Unexpected type {type(tensor)}"
-        assert tensor is None or tensor.numel() > 1 or tensor.item() != -666666
+        assert isinstance(tensor, torch.Tensor), (
+            f"Unexpected type {type(tensor)}, "
+            f"__name__={getattr(tensor, '__name__', 'no name')}"
+        )
+        assert tensor is None or tensor.numel() != 1 or tensor.item() != -666666
         self.tensor = (
             tensor.cpu()
             if may_cpu
@@ -197,6 +200,7 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
             ), f"Cannot guess version from name={self.__class__.__name__!r}"
             version = int(name[1])
         self.version = version
+        self.name = node.name
 
     def __str__(self) -> str:
         "usual"
@@ -291,6 +295,15 @@ def get_attribute_tensor(self, node: onnx.NodeProto, name: str) -> Optional[torc
             return None
         return to_tensor(att.t)
 
+    def same_device(self, *tensors: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        """Puts all tensors on the same device."""
+        devices = [t.get_device() for t in tensors]
+        if len(set(devices)) == 1:
+            return tuple(tensors)
+        index = devices.index(max(devices))
+        device = tensors[index].device
+        return tuple(t.to(device) for t in tensors)
+
 
 class OpRunFunction(OpRun):
     """

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -65,11 +65,30 @@ def run(self, x, w, b=None):
         ), f"conv not implemented for kernel_shape={kernel_shape} and w.shape={w.shape}"
         dilations = self.dilations or [1 for _ in x.shape[2:]]
         strides = self.strides or [1 for _ in x.shape[2:]]
-        pads = self.pads or ([0 for _ in x.shape[2:]] * 2)
-        assert (
-            self.auto_pad == "NOTSET"
-        ), f"conv not implemented for auto_pad={self.auto_pad!r}"
-        assert len(set(pads)) == 1, f"conv not implemented for pads={pads}"
+
+        if self.auto_pad in {"SAME_LOWER", "SAME_UPPER"}:
+            head = []
+            tail = []
+            for i in range(len(x.shape) - 2):
+                d = x.shape[i + 2]
+                target_size = (d + strides[i] - 1) // strides[i]
+                pad_needed = (target_size - 1) * strides[i] + kernel_shape[i] - d
+                pad_head = (
+                    (pad_needed + 1) // 2 if self.auto_pad == "SAME_LOWER" else pad_needed // 2
+                )
+                pad_tail = pad_needed - pad_head
+                head.append(pad_head)
+                tail.append(pad_tail)
+            pads = head + tail
+        else:
+            pads = self.pads or ([0 for _ in x.shape[2:]] * 2)
+
+        assert len(set(pads)) == 1, (
+            f"conv not implemented for pads={pads}, "
+            f"auto_pad={self.auto_pad!r}, strides={strides}, "
+            f"x.shape={x.shape}, kernel_shape={kernel_shape}"
+        )
+
         if b is None:
             bias = None
         else:

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -42,7 +42,21 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         self.axis = axis
 
     def run(self, *data: OpRunTensor) -> OpRunTensor:
-        return OpRunTensor(torch.cat([t.tensor for t in data], axis=self.axis))
+        assert data, f"No tensor to concatenate in node name {self.name!r}"
+        devices = [d.get_device() for d in data]
+        if len(set(devices)) == 1:
+            return OpRunTensor(torch.cat([t.tensor for t in data], axis=self.axis))
+        if (
+            data[0].dtype == torch.int64
+            and self.axis == 0
+            and max(d.tensor.ndim for d in data) == 1
+            and max(d.tensor.numel() for d in data) <= 8
+        ):
+            # This is a shape
+            return OpRunTensor(torch.cat([t.tensor.cpu() for t in data], axis=self.axis))
+        index = devices.index(max(devices))
+        device = data[index].tensor.device
+        return OpRunTensor(torch.cat([t.tensor.to(device) for t in data], axis=self.axis))
 
 
 class NonZero_13(OpRun):
@@ -88,4 +102,5 @@ class Where_9(OpRun):
     "Where"
 
     def run(self, cond: OpRunTensor, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
-        return OpRunTensor(torch.where(cond.tensor, x.tensor, y.tensor))
+        tcond, tx, ty = self.same_device(cond.tensor, x.tensor, y.tensor)
+        return OpRunTensor(torch.where(tcond, tx, ty))

onnx_diagnostic/reference/torch_ops/reduce_ops.py

@@ -42,7 +42,7 @@ def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor
             assert (
                 not self.keepdims
             ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
-            return OpRunTensor(torch.max(x.tensor).values)
+            return OpRunTensor(x.tensor.max())
         taxes = axes.as_tuple_int
         if len(taxes) == 1:
             t = x.tensor.max(taxes[0], keepdim=self.keepdims)
@@ -81,7 +81,7 @@ def run(self, x: OpRunTensor) -> OpRunTensor:
             assert (
                 not self.keepdims
             ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
-            return OpRunTensor(torch.min(x.tensor).values)
+            return OpRunTensor(x.tensor.min())
         taxes = tuple(axes)
         if len(taxes) == 1:
             t = x.tensor.min(taxes[0], keepdim=self.keepdims)
@@ -100,8 +100,8 @@ def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor
         if axes is None:
             assert (
                 not self.keepdims
-            ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
-            return OpRunTensor(torch.min(x.tensor).values)
+            ), f"axes is empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(torch.min(x.tensor))
         taxes = axes.as_tuple_int
         if len(taxes) == 1:
             t = x.tensor.min(taxes[0], keepdim=self.keepdims)