From 4c07ef278632ea3831331c4506aa49c47fe8dea6 Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Fri, 30 May 2025 10:16:04 -0700 Subject: [PATCH] [ET-VK] Modifying should_squeeze function in SqueezeUnsqueezeInputs to not squeeze if significant axis are all 1 and trailing axis are all > 1. Pull Request resolved: https://github.com/pytorch/executorch/pull/11177 This diff modifies the `should_squeeze` function in `SqueezeUnsqueezeInputs` to not squeeze (return False) if significant axes are all 1 and trailing axes are all > 1. ghstack-source-id: 287222796 @exported-using-ghexport Differential Revision: [D75483587](https://our.internmc.facebook.com/intern/diff/D75483587/) --- backends/vulkan/_passes/squeeze_unsqueeze_inputs.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py b/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py index b4337829d7f..c415249383e 100644 --- a/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py +++ b/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py @@ -32,7 +32,13 @@ def should_squeeze(self, op, shape: List[int]) -> bool: # pyre-ignore return shape[1] == 1 and shape[0] > 1 if len(shape) == 4: # No need to squeeze if all dims are 1 except the width dim - if all(dim == 1 for dim in shape[:-1]): + if shape[0] == shape[1] == shape[2] == 1: + return False + # No need to squeeze if batch and channel dims are 1 and height and width are > 1 + if shape[0] == shape[1] == 1 and shape[2] > 1 and shape[3] > 1: + return False + # No need to squeeze if batch dim is 1 and channel, height and width are > 1 + if shape[0] == 1 and shape[1] > 1 and shape[2] > 1 and shape[3] > 1: return False # Otherwise, check for squeezable dim return 1 in shape[:-1]