🐛 Bug

torch.linalg.qr does not work correctly when input is a batched tensor with some of the batch dimensions set to zero.

To Reproduce

In [1]: import torch
In [2]: shape = (1, 0, 2, 4, 5)
In [3]: A = torch.randn(shape)
In [4]: Q, R = torch.linalg.qr(A)
In [5]: Q.shape
Out[5]: torch.Size([4, 4])
In [6]: R.shape
Out[6]: torch.Size([4, 5])

Expected behavior

Q should have the shape (1, 0, 2, 4, 4) and R should have the shape (1, 0, 2, 4, 5) for the above example.

cc @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @IvanYashchuk @antocuni