Make squeeze and slice_by_size const folding agree with their type inference

coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py

@@ -670,7 +670,9 @@ def value_inference(self):
                 if is_symbolic(self.x.shape[i]):
                     return None
                 begin_val += self.x.shape[i]
-            if self.size.val[i] > 0:
+            if self.size.val[i] != -1:
+                # As in type_inference, -1 is the only value that means "the rest of
+                # the dimension"; every other size is taken literally.
                 slices.append(slice(begin_val, begin_val + self.size.val[i]))
             else:
                 slices.append(slice(begin_val, None, None))
@@ -941,8 +943,13 @@ def value_inference(self):
         if self.axes is None:
             val = np.squeeze(self.x.val)
         else:
-            val = np.squeeze(self.x.val, axis=tuple(self.axes.val))
-        return val if val.shape != () else self.x.val[0]
+            axes = [axis if axis >= 0 else axis + self.x.rank for axis in self.axes.val]
+            # As in type_inference, an axis whose size is not 1 is ignored rather than
+            # raising, which is the PyTorch behavior this op documents.
+            axes = tuple(axis for axis in axes if self.x.val.shape[axis] == 1)
+            val = np.squeeze(self.x.val, axis=axes)
+        # A zero rank result must be returned as a scalar, not as a zero rank array.
+        return val if val.shape != () else val[()]
 
 @register_op
 class transpose(Operation):

coremltools/converters/mil/mil/ops/tests/iOS14/test_tensor_transformation.py

@@ -1210,6 +1210,17 @@ def test_builder_eval(self):
         np.testing.assert_allclose(x[:, 1:, :3], v_2.val, atol=1e-04, rtol=1e-05)
         np.testing.assert_allclose(x[:, -2:, :3], v_3.val, atol=1e-04, rtol=1e-05)
 
+    @ssa_fn
+    def test_builder_eval_zero_size(self):
+        """
+        Only -1 means "the rest of the dimension"; a 0 size is an empty slice, which is
+        what the type inference already reports.
+        """
+        x = np.array(list(range(24))).reshape(2, 3, 4)
+        v = mb.slice_by_size(x=x, begin=(0, 1, 0), size=(-1, 0, 3))
+        assert v.shape == (2, 0, 3)
+        assert v.val.shape == (2, 0, 3)
+
 
 class TestSpaceToDepth:
     @pytest.mark.parametrize(
@@ -1302,6 +1313,34 @@ def test_builder_eval_rank_0(self):
         assert type(v.val) == np.float32
         assert np.isclose(np.squeeze(x), v.val)
 
+    @ssa_fn
+    def test_builder_eval_rank_0_from_higher_rank(self):
+        """A zero rank result must be a scalar, whatever the rank of the input was."""
+        for shape in [(1,), (1, 1), (1, 1, 1)]:
+            x = np.full(shape, 5.0, dtype=np.float32)
+            v = mb.squeeze(x=x)
+            assert v.shape == ()
+            assert type(v.val) == np.float32
+            assert np.isclose(np.squeeze(x), v.val)
+
+    @ssa_fn
+    def test_builder_eval_non_single_element_dim(self):
+        """
+        The const folded value must follow the same "ignore non single dimensions"
+        rule that the type inference and the runtime follow.
+        """
+        x = np.arange(2 * 3 * 4, dtype=np.int32).reshape(2, 3, 4)
+        for axes in [(-1,), (-2, 0), (0, 1, 2)]:
+            v = mb.squeeze(x=x, axes=axes)
+            assert v.shape == x.shape
+            np.testing.assert_array_equal(x, v.val)
+
+        # Mixing squeezable and non squeezable axes only drops the squeezable ones.
+        y = np.arange(2 * 3, dtype=np.int32).reshape(1, 2, 1, 3)
+        v = mb.squeeze(x=y, axes=(0, 1, 2))
+        assert v.shape == (2, 3)
+        np.testing.assert_array_equal(y.reshape(2, 3), v.val)
+
     @staticmethod
     def test_squeeze_value_inference_is_inplace():
         @mb.program()