Skip to content

Make squeeze and slice_by_size const folding agree with their type inference - #2803

Open
LeSingh1 wants to merge 1 commit into
apple:mainfrom
LeSingh1:tensor-transformation-value-inference
Open

Make squeeze and slice_by_size const folding agree with their type inference#2803
LeSingh1 wants to merge 1 commit into
apple:mainfrom
LeSingh1:tensor-transformation-value-inference

Conversation

@LeSingh1

@LeSingh1 LeSingh1 commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

Three places in iOS15/tensor_transformation.py where value_inference disagrees with the type_inference (and the docstring) of the same op, so a program behaves differently depending only on whether its input happens to be const.

squeeze raises on a non-single dimension

The docstring says, verbatim:

The behaviour of squeezing non-single dimensions follow PyTorch instead of NumPy, where it ignores non-single dimensions instead of erroring out. More specifically, if x has shape (2, 3, 4) and axes is [0, 1], the output will be a tensor with shape (2, 3, 4).

type_inference implements exactly that (if squeezed_shape[i] == 1:), and TestSqueeze::test_non_single_element_dim pins the runtime behavior. value_inference passed the axes straight to np.squeeze, which is the NumPy rule:

>>> mb.squeeze(x=np.arange(24, dtype=np.float32).reshape(2, 3, 4), axes=[0, 1])
ValueError: cannot select an axis to squeeze out which has size not equal to one

That is the docstring's own example, and it aborts the conversion. The same program converts and runs when x is a model input rather than a const.

squeeze returns a rank 1 array where the type says scalar

return val if val.shape != () else self.x.val[0]

self.x.val[0] is a scalar only when x is rank 1. For an all-ones shape of rank 2 or more it is still an array, so building the op fails:

>>> mb.squeeze(x=np.array([[5.0]], dtype=np.float32))
ValueError: Types should have zero-rank ndarray input, got [5.] instead.

Fixed by indexing the squeezed value with (), which gives the same numpy scalar the rank 1 case already produced.

slice_by_size treats a size of 0 as "the rest of the dimension"

if self.size.val[i] > 0:
    slices.append(slice(begin_val, begin_val + self.size.val[i]))
else:
    slices.append(slice(begin_val, None, None))

The docstring gives that meaning to -1 only, and type_inference agrees (elif s != -1: ret_shape.append(s)). With size = [0, 2] on a (3, 4) input and begin = [1, 1], type_inference reports shape (0, 2) while value_inference returns a (2, 2) array — a const whose value contradicts its own declared type. Changed the guard to != -1 so the two agree.

Testing

Three new builder-eval tests, all of which fail on main:

  • TestSqueeze::test_builder_eval_non_single_element_dim uses the same axes as the existing test_non_single_element_dim backend test, so the folded value is pinned to the runtime value.
  • TestSqueeze::test_builder_eval_rank_0_from_higher_rank extends the existing test_builder_eval_rank_0 to ranks 2 and 3.
  • TestSliceBySize::test_builder_eval_zero_size checks the folded shape against the declared shape.

The existing TestSqueeze / TestSliceBySize tests are unchanged and still pass, including test_squeeze_value_inference_is_inplace (the folded value is still a view of the const) and test_builder_eval_rank_0. The full iOS14/16/17/18 test_tensor_transformation.py suites pass.

…ference

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant