Can't solve torch.lstsq() with specific values

[xiaosu-zhu]

🐛 Bug

When using torch.lstsq to solve with specific values, torch.lstsq produces wrong result (on cuda) or raises error (on cpu) while scipy.linalg.lstsq not.

To Reproduce

Steps to reproduce the behavior:

1. Use the dumped .npy file: a.npy, b.npy to perform torch.lstsq
2. Meanwhile, lstsq can also be solved in scipy.linalg.lstsq
3. With the given numpy array, torch.lstsq raises error on cpu, or gives wrong result on cuda, while scipy.linalg.lstsq not.

Code samples and error messages:

Download files:

• a.npy: https://github.com/xiaosu-zhu/archives/releases/download/torch.lstsq/a.npy
• b.npy: https://github.com/xiaosu-zhu/archives/releases/download/torch.lstsq/b.npy

import numpy as np
import scipy.linalg
import torch

# sparse 0-1 matrix with shape [100000, 1024]
# each row contains exactly 4 ones and other zeros
a = np.load("a.npy").astype(np.float32)
# int matrix with shape [100000, 128]
b = np.load("b.npy").astype(np.float32)

x, _, _, _ = scipy.linalg.lstsq(a, b)

# 67480.16036171981
print(((a @ x - b) ** 2).sum(-1).mean())

b = torch.from_numpy(b).cuda()
a = torch.from_numpy(a).cuda()
x, _ = torch.lstsq(b, a)

# 7.4439e+15
print(((a @ x[:a.shape[-1]] - b) ** 2).sum(-1).mean())

# raise RuntimeError: Lapack Error in gels : 
# The 259-th diagonal element of the triangular factor of A is zero 
# at /opt/conda/conda-bld/pytorch_1595629395347/work/aten/src/TH/generic/THTensorLapack.cpp:177
x, _ = torch.lstsq(b.cpu(), a.cpu())

Raise RuntimeError: Lapack Error in gels : The 259-th diagonal element of the triangular factor of A is zero at /opt/conda/conda-bld/pytorch_1595629395347/work/aten/src/TH/generic/THTensorLapack.cpp:177 when using torch.FloatTensor.
Or an irregular result when using torch.cuda.FloatTensor.

Expected behavior

Produce the same result with scipy.linalg.lstsq.

Environment

PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-10ubuntu2) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti
GPU 2: GeForce GTX 1080 Ti
GPU 3: GeForce GTX 1080 Ti
GPU 4: GeForce GTX 1080 Ti
GPU 5: GeForce GTX 1080 Ti
GPU 6: GeForce GTX 1080 Ti
GPU 7: GeForce GTX 1080 Ti

Nvidia driver version: 450.51.05
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               10.2.89              hfd86e86_1    defaults
[conda] mkl                       2020.1                      217    defaults
[conda] mkl-service               2.3.0            py38he904b0f_0    defaults
[conda] mkl_fft                   1.1.0            py38h23d657b_0    defaults
[conda] mkl_random                1.1.1            py38h0573a6f_0    defaults
[conda] numpy                     1.19.1           py38hbc911f0_0    defaults
[conda] numpy-base                1.19.1           py38hfa32c7d_0    defaults
[conda] pytorch                   1.6.0           py3.8_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] torchvision               0.7.0                py38_cu102    pytorch

cc @ezyang @gchanan @zou3519 @bdhirsh @vishwakftw @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @rgommers @vincentqb

[mrshenli]

Hey @mruberry, do we also have plan to add a lstsq to torch.linalg and retire existing ones in the near future?

[mruberry]

Yes we do, see #42666. We would definitely take a fix for this.

[xiaosu-zhu]

Found related issue #11174

[chrisby]

In fact, I stumbled over the fact that you can reproduce this very easily with a 2 x 2 matrix:

from scipy import linalg
import torch

a = np.array([[0., 0.], [0., 1.]])
b = np.array([[0., 0.], [0., 1.]])

linalg.lstsq(a,b)[0] # Yields [[0., 0.], [0., 1.]]

a_t = torch.from_numpy(a)
b_t = torch.from_numpy(b)
torch.lstsq(a_t,b_t)

The last line throws RuntimeError: Lapack Error in gels : The 1-th diagonal element of the triangular factor of A is zero at ../aten/src/TH/generic/THTensorLapack.cpp:177

Is there a timeline on a fix? Seems to be rather relevant.

[mruberry]

Thanks for the simpler reproduction, @chrisby.

Our team is currently working through linear algebra updates, but I've changed this issue to high priority. It should be fixed before the next release.

[xiaosu-zhu]

Thank you, @chrisby @mruberry. Look forward to the next release ❤️