Fixed division by zero in QR decomposition. Issue #1058

src/ops/linalg_ops.ts

@@ -211,7 +211,12 @@ function qr2d(x: Tensor2D, fullMatrices = false): [Tensor2D, Tensor2D] {
         const rjEnd1 = r.slice([j, j], [m - j, 1]);
         const normX = rjEnd1.norm();
         const rjj = r.slice([j, j], [1, 1]);
-        const s = rjj.sign().neg() as Tensor2D;
+
+        // The sign() function returns 0 on 0, which causes division by zero.
+        const s = tensor2d([[-1]]).where(
+            rjj.greater(tensor2d([[0]])),
+            tensor2d([[1]]));
+
         const u1 = rjj.sub(s.mul(normX)) as Tensor2D;
         const wPre = rjEnd1.div(u1);
         if (wPre.shape[0] === 1) {

src/ops/linalg_ops_test.ts

@@ -171,6 +171,29 @@ describeWithFlags('qr', ALL_ENVS, () => {
             [3, 3]));
   });
 
+  it('3x3, zero on diagonal', () => {
+    const x = tensor2d([[0, 2, 2], [1, 1, 1], [0, 1, 2]], [3, 3]);
+    const [q, r] = tf.linalg.qr(x);
+    expectArraysClose(
+        q,
+        tensor2d(
+            [
+              [0., -0.89442719, 0.4472136],
+              [1., 0., 0.],
+              [0., -0.4472136, -0.89442719]
+            ],
+            [3, 3]));
+    expectArraysClose(
+        r,
+        tensor2d(
+            [
+              [1., 1., 1.],
+              [0., -2.23606798, -2.68328157],
+              [0., 0., -0.89442719]
+            ],
+            [3, 3]));
+  });
+
   it('3x2, fullMatrices = default false', () => {
     const x = tensor2d([[1, 2], [3, -3], [-2, 1]], [3, 2]);
     const [q, r] = tf.linalg.qr(x);