Swin_v2_t fails at runtime on Vulkan

🐛 Describe the bug

The swin_v2_t model from torchvision fails at runtime on the Vulkan backend with "RuntimeError: Exception raised from get_uniform_data at /Users/gjcomer/src/executorch/../executorch/backends/vulkan/runtime/api/containers/Tensor.h:679: (sizes_.size() <= 4) is false!".

This can be reproduced with the following test case command or standalone script.

python -m executorch.backends.test.suite.runner models --flow vulkan --filter "test_swin_v2_t_vulkan_float32$"

Standalone repro:

import torch
import torchvision
from executorch.exir import to_edge_transform_and_lower

from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner

inputs = (torch.randn(1, 3, 224, 224),)
model = torchvision.models.swin_v2_t().eval()

ep = torch.export.export(model, inputs)

model = to_edge_transform_and_lower(
    torch.export.export(model, inputs),
    partitioner=[VulkanPartitioner()],
).to_executorch()

print("Running model...")
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
loaded_model = _load_for_executorch_from_buffer(model.buffer)
loaded_model([*inputs])

Note that running the backend test case requires executorch's python bindings to be built with the Vulkan backend. An example build command is below.

CMAKE_ARGS="-DEXECUTORCH_BUILD_VULKAN=ON" ./install_executorch.sh --editable

Versions

Commit fbda3a9545de747329577bd910086072ec5c7ad1, M1 Mac, using MoltenVK

cc @SS-JIA @manuelcandales @cbilgin

Build the Python bindings with Vulkan enabled, then run the listed test command or standalone Swin v2_t reproduction. Start at the reported Vulkan runtime location, executorch/backends/vulkan/runtime/api/containers/Tensor.h:679, and trace the failing tensor shape. Done means the test model loads and runs on Vulkan without the runtime error.