diff --git a/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py b/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py index 67fdf23d8..b0badfe5b 100644 --- a/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py +++ b/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): diff --git a/coremltools/converters/mil/mil/ops/tests/iOS14/test_tensor_transformation.py b/coremltools/converters/mil/mil/ops/tests/iOS14/test_tensor_transformation.py index 5dd4495ba..134b4b901 100644 --- a/coremltools/converters/mil/mil/ops/tests/iOS14/test_tensor_transformation.py +++ b/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()