From 79e0dc9140e132c485cfaac39a77cab54dce817e Mon Sep 17 00:00:00 2001 From: LeSingh1 Date: Sun, 9 Aug 2026 13:29:31 -0700 Subject: [PATCH] Make squeeze and slice_by_size const folding agree with their type inference squeeze documents, and type_inference implements, PyTorch's rule that an axis whose size is not 1 is ignored instead of raising. value_inference handed the axes straight to np.squeeze, which raises. So squeezing a const along a non-single dimension aborted the conversion, even though the same program converts and runs when the input is not const. squeeze's zero rank result was returned as self.x.val[0], which is only a scalar when x is rank 1. For an all-ones shape of rank 2 or more this returned an array, and building the op failed with "Types should have zero rank ndarray input". Index the squeezed value instead. slice_by_size's value_inference treated any size <= 0 as "the rest of the dimension", while its type_inference (and the docstring) give that meaning to -1 only. A size of 0 therefore produced a const whose value disagreed with its own declared shape. --- .../ops/defs/iOS15/tensor_transformation.py | 13 +++++-- .../tests/iOS14/test_tensor_transformation.py | 39 +++++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) 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()