Output shape of torch.linalg.norm diverges from numpy equivalent in edge case.

[pmeier]

🐛 Bug

Output shape of torch.linalg.norm diverges from numpy equivalent in edge case.

To Reproduce

import torch
import numpy as np

t = torch.rand((1, 1, 0))
a = t.numpy()

t_norm = torch.linalg.norm(t, ord=2, dim=(0, 1), keepdim=False)
a_norm = np.linalg.norm(a, ord=2, axis=(0, 1), keepdims=False)

assert t_norm.shape == a_norm.shape

• This also fails for ord=-2 but works for all other orders
• This works fine for keepdim=True
• This also works fine for other degenerate input shapes

Environment

PyTorch version: 1.9.0a0+gitc9de3e3
Is debug build: True
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (crosstool-NG 1.24.0.133_b0863d8_dirty) 9.3.0
Clang version: Could not collect
CMake version: version 3.20.2
Libc version: glibc-2.9

Python version: 3.6 (64-bit runtime)
Python platform: Linux-5.12.5-arch1-1-x86_64-with-arch
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1080
Nvidia driver version: 465.31
cuDNN version: Probably one of the following:
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn.so.8.1.1
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn_adv_infer.so.8.1.1
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn_adv_train.so.8.1.1
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn_cnn_infer.so.8.1.1
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn_cnn_train.so.8.1.1
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn_ops_infer.so.8.1.1
/usr/local/cudnn-8.1.1-cuda-11/lib64/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] pytest-pytorch==0.0.0
[pip3] torch==1.9.0a0+gitc9de3e3
[conda] magma-cuda111             2.5.2                         1    pytorch
[conda] mkl                       2021.2.0           h726a3e6_389    conda-forge
[conda] mkl-include               2021.2.0           h726a3e6_389    conda-forge
[conda] numpy                     1.19.5           py36h2aa4a07_1    conda-forge
[conda] pytest-pytorch            0.2.0              pyh44b312d_0    conda-forge
[conda] torch                     1.9.0a0+gitc9de3e3           dev_0    <develop>

Additional context

This is tested in

pytorch/test/test_linalg.py

Line 1928 in 6d45d7a

def test_norm_matrix_degenerate_shapes(self, device, dtype):

but the test never failed due to a bug in the underlying comparison mechanism. #59067 will disable the failing test case, but it should be reinstated as soon as this issue is fixed.

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @IvanYashchuk @xwang233 @lezcano @VitalyFedyunin

[mruberry]

Thank you for identifying this issue, @pmeier. We would definitely accept a fix for this!