Skip to content

[ET-VK] Fix squeeze_copy of the outermost dim under dynamic shapes - #22409

Open
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/vulkan-squeeze-dim0-resize
Open

[ET-VK] Fix squeeze_copy of the outermost dim under dynamic shapes#22409
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/vulkan-squeeze-dim0-resize

Conversation

@msluszniak

Copy link
Copy Markdown
Contributor

Summary

add_squeeze_copy_dims_node() deliberately skips dim 0 and falls back to add_clone_node(). But resize_clone_node() only propagates sizes when input and output have the same dim count -- which a squeeze never does, and the code says so:

// TODO: support for when dimensionality doesn't match, i.e. clone is used to
// implement squeeze.
if (graph->dim_of(out) == graph->dim_of(in)) {
  graph->virtual_resize(out, graph->sizes_of(in));
}

So the output is never resized. With static shapes that is invisible. With dynamic shapes the output keeps its upper-bound extents while consumers read it at the real size, the copy lands in the wrong places, and roughly half the output comes back zeroed. There is no error; the values are just wrong.

op_registry.py declares supports_resize=True for squeeze_copy.dims, so the partitioner accepts the op and the result is a silent wrong answer rather than a rejected partition.

Fix

Route dim 0 through the permute path like every other squeeze dim. resize_permute_node() already has an explicit branch for the rank-reducing case, so no new resize logic is needed.

Reproduction

Any model that ends up with torch.cat(list(x), -1) over a rank-4 tensor with a dynamic dim: the unbind lowers to slice_copy + squeeze_copy.dims, and the second slice comes back zeroed for every extent below the bound. Reduced:

class M(torch.nn.Module):
    def forward(self, y):
        return y[1:2], y[1:2].squeeze(0)

ep = torch.export.export(M(), (torch.randn(2, 2, 1000, 128),),
                         dynamic_shapes=({2: Dim("L", min=8, max=1000)},))

At L=200, y[1:2] is correct while y[1:2].squeeze(0) returns exactly half zeros -- cosine 0.704 against the reference, which is sqrt(0.5). It is correct only at L=1000, the bound.

Verification

Galaxy S26 Ultra, Adreno 840, fp16:

case before after
y[1:2].squeeze(0), L=200 / 500 / 1000 0.704 / 0.708 / 1.000000 1.000000 at all three
TTS model, CFG batch built via cat(list(x), -1) 0.36 0.99993

The second row is a supertonic TTS vector-estimator whose classifier-free-guidance uncond branch was the zeroed slice; it is now within fp16 noise of its CPU reference at every sequence length tested (64, 200, 500, 1000).

Note that this class of bug is invisible to executor_runner, which can only run a model at its dynamic upper bound -- there is no flag to request a smaller shape, and --inputs files must match nbytes() of the bound. I found it with a local patch adding an --input_shapes flag; happy to send that separately if it would be useful.

add_squeeze_copy_dims_node() skips dim 0 and falls back to
add_clone_node(). resize_clone_node() only propagates sizes when input
and output have the same dim count, which a squeeze never does, so the
output keeps the extents it was built with.

With static shapes that is invisible. With dynamic shapes the output
holds its upper-bound extents while consumers read it at the real size,
so the copy lands in the wrong places and roughly half the output comes
back zeroed -- silently, with no error.

Route dim 0 through the permute path like every other squeeze dim;
resize_permute_node() already has an explicit branch for the
rank-reducing case.

Repro: any model that ends up with torch.cat(list(x), -1) over a rank-4
tensor with a dynamic dim. The unbind lowers to slice_copy plus
squeeze_copy.dims, and the second slice comes back zeroed for every
extent below the bound. Reduced to a 15-line case: y[1:2] is correct
while y[1:2].squeeze(0) returns exactly half zeros (cosine 0.704 = sqrt
of 0.5 against the reference), correct only at the bound.

Verified on a Galaxy S26 Ultra (Adreno 840): the reduced case goes from
0.704 to 1.000000 at extents 200, 500 and 1000, and a TTS model whose
classifier-free-guidance batch is built this way goes from cosine 0.36
to 0.99993 against its CPU reference.
@msluszniak
msluszniak requested a review from SS-JIA as a code owner September 1, 2026 13:44
@pytorch-bot

pytorch-bot Bot commented Sep 1, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22409

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit a4629cf with merge base 39d623e (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 1, 2026
@github-actions

github-actions Bot commented Sep 1, 2026

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants