pytorch - 💡(How to fix) Fix torch.compile returns wrong shape for torch.norm(dim=-1) when reduced dimension has size 1 [1 participants]

🐛 Describe the bug

torch.compile with the Inductor backend returns an incorrect output shape for torch.norm(dim=-1) when the reduced last dimension has size 1.

Eager execution correctly removes the reduced dimension and returns shape (5, 7), while the compiled function returns shape (5, 7, 1).

import torch
import torch.nn as nn

class M(nn.Module):
    def forward(self, x):
        return x.norm(dim=-1)

m = M()
torch._dynamo.reset()
c = torch.compile(m)

x = torch.rand(5, 7, 1)

es = m(x).shape
cs = c(x).shape

print(f"eager   shape = {tuple(es)}")
print(f"compile shape = {tuple(cs)}")
print("BUG" if es != cs else "OK")

Error logs

eager shape = (5, 7) compile shape = (5, 7, 1) BUG

Versions

PyTorch version: 2.11.0+cu130 Is debug build: False CUDA used to build PyTorch: 13.0 ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.2 LTS (x86_64) GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0 Clang version: Could not collect CMake version: Could not collect Libc version: glibc-2.39

Python version: 3.10.20 (main, Mar 11 2026, 17:46:40) [GCC 14.3.0] (64-bit runtime) Python platform: Linux-5.15.167.4-microsoft-standard-WSL2-x86_64-with-glibc2.39 Is CUDA available: False CUDA runtime version: Could not collect CUDA_MODULE_LOADING set to: N/A GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080 Laptop GPU Nvidia driver version: 545.92 cuDNN version: Could not collect Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

Versions of relevant libraries: [pip3] numpy==2.2.6 [pip3] onnx==1.21.0 [pip3] onnx2torch==1.5.15 [pip3] onnxruntime==1.23.2 [pip3] torch==2.11.0 [pip3] torchvision==0.26.0 [pip3] triton==3.6.0

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo