Repository: mars-project/mars

Issue:
[BUG] Tensors returned by QR decomposition have wrong nsplits after tile
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Tensors returned by QR decomposition have wrong `nsplits` after tile

**To Reproduce**

```
In [62]: q, r = mt.linalg.qr(a)                                                 

In [63]: q.tiles()                                                              
Out[63]: Tensor <op=TensorQR, shape=(10, 4), key=d0eb552be4c832153edf2e245410dc44>

In [64]: r.tiles()                                                              
Out[64]: Tensor <op=TensorQR, shape=(4, 4), key=18580529a6db43e1789f4c7ee721d425>

In [65]: q.nsplits                                                              
Out[65]: ((1,), (1,))  # wrong, should be ((10,), (4,))

In [66]: q.shape                                                                
Out[66]: (10, 4)

In [67]: r.nsplits                                                              
Out[67]: ((1,), (1,))  # wrong, should be ((10,), (4,))

In [68]: r.shape                                                                
Out[68]: (4, 4)
```




Discussion:

Resolving patch:
diff --git a/mars/core.py b/mars/core.py
--- a/mars/core.py
+++ b/mars/core.py
@@ -582,7 +582,7 @@ def _new_tileables(self, inputs, kws=None, **kw):
             # for each output tileable, hold the reference to the other outputs
             # so that either no one or everyone are gc collected
             for j, t in enumerate(tileables):
-                t.data._siblings = [t.data for t in tileables[:j] + tileables[j + 1:]]
+                t.data._siblings = [tileable.data for tileable in tileables[:j] + tileables[j + 1:]]
         return tileables
 
     def new_tileables(self, inputs, kws=None, **kw):
diff --git a/mars/tensor/expressions/linalg/core.py b/mars/tensor/expressions/linalg/core.py
--- a/mars/tensor/expressions/linalg/core.py
+++ b/mars/tensor/expressions/linalg/core.py
@@ -71,7 +71,7 @@ def tile(cls, op):
         q, r = op.outputs
         new_op = op.copy()
         q_nsplits = ((q_chunk.shape[0],), (q_chunk.shape[1],))
-        r_nsplits = ((1,), (c.shape[1] for c in r_chunks))
+        r_nsplits = ((r_chunks[0].shape[0],), tuple(c.shape[1] for c in r_chunks))
         kws = [
             {'chunks': [q_chunk], 'nsplits': q_nsplits, 'dtype': q.dtype, 'shape': q.shape},
             {'chunks': r_chunks, 'nsplits': r_nsplits, 'dtype': r.dtype, 'shape': r.shape}
@@ -121,9 +121,9 @@ def tile(cls, op):
         qr_op = TensorQR()
         qr_chunks = qr_op.new_chunks([concat_r_chunk], index=concat_r_chunk.index,
                                      kws=[{'side': 'q', 'dtype': q_dtype,
-                                           'shape': concat_r_chunk.shape},
+                                           'shape': (concat_r_chunk.shape[0], min(concat_r_chunk.shape))},
                                           {'side': 'r', 'dtype': r_dtype,
-                                           'shape': (concat_r_chunk.shape[1],) * 2}])
+                                           'shape': (min(concat_r_chunk.shape), concat_r_chunk.shape[1])}])
         stage2_q_chunk, stage2_r_chunk = qr_chunks
 
         # stage 3, map phase
@@ -134,7 +134,8 @@ def tile(cls, op):
         stage2_q_chunks = []
         for c, s in zip(stage1_q_chunks, q_slices):
             slice_op = TensorSlice(slices=[s], dtype=c.dtype)
-            stage2_q_chunks.append(slice_op.new_chunk([stage2_q_chunk], shape=c.shape, index=c.index))
+            stage2_q_chunks.append(slice_op.new_chunk([stage2_q_chunk], index=c.index,
+                                                      shape=(c.shape[0], stage2_q_chunk.shape[1])))
         stage3_q_chunks = []
         for c1, c2 in izip(stage1_q_chunks, stage2_q_chunks):
             dot_op = TensorDot(dtype=q_dtype)
@@ -144,8 +145,7 @@ def tile(cls, op):
         if not calc_svd:
             q, r = op.outputs
             new_op = op.copy()
-            # unify_nsplits will get nsplits by chunk shape if it was set to (1,)
-            q_nsplits = ((c.shape[0] for c in stage3_q_chunks), (1,))
+            q_nsplits = (tuple(c.shape[0] for c in stage3_q_chunks), (stage3_q_chunks[0].shape[1],))
             r_nsplits = ((stage2_r_chunk.shape[0],), (stage2_r_chunk.shape[1],))
             kws = [
                 # Q
@@ -185,7 +185,7 @@ def tile(cls, op):
                                                             index=c1.index))
 
             new_op = op.copy()
-            u_nsplits = ((c.shape[0] for c in stage4_u_chunks), (1,))
+            u_nsplits = (tuple(c.shape[0] for c in stage4_u_chunks), (stage4_u_chunks[0].shape[1],))
             s_nsplits = ((stage2_s_chunk.shape[0],),)
             v_nsplits = ((stage2_v_chunk.shape[0],), (stage2_v_chunk.shape[1],))
             kws = [
diff --git a/mars/tensor/expressions/linalg/qr.py b/mars/tensor/expressions/linalg/qr.py
--- a/mars/tensor/expressions/linalg/qr.py
+++ b/mars/tensor/expressions/linalg/qr.py
@@ -68,8 +68,10 @@ def tile(cls, op):
 
             new_op = op.copy()
             kws = [
-                {'chunks': [q_chunk], 'nsplits': ((1,), (1,)), 'dtype': q_dtype, 'shape': q_shape},
-                {'chunks': [r_chunk], 'nsplits': ((1,), (1,)), 'dtype': r_dtype, 'shape': r_shape}
+                {'chunks': [q_chunk], 'nsplits': ((q_shape[0],), (q_shape[1],)),
+                 'dtype': q_dtype, 'shape': q_shape},
+                {'chunks': [r_chunk], 'nsplits': ((r_shape[0],), (r_shape[1],)),
+                 'dtype': r_dtype, 'shape': r_shape}
             ]
             return new_op.new_tensors(op.inputs, kws=kws)
         elif op.method == 'tsqr':


Test patch:
diff --git a/mars/tensor/expressions/tests/test_linalg.py b/mars/tensor/expressions/tests/test_linalg.py
--- a/mars/tensor/expressions/tests/test_linalg.py
+++ b/mars/tensor/expressions/tests/test_linalg.py
@@ -39,6 +39,8 @@ def testQR(self):
 
         self.assertEqual(len(q.chunks), 3)
         self.assertEqual(len(r.chunks), 1)
+        self.assertEqual(q.nsplits, ((3, 3, 3), (6,)))
+        self.assertEqual(r.nsplits, ((6,), (6,)))
         self.assertEqual(calc_shape(q.chunks[0]), q.chunks[0].shape)
         self.assertEqual(calc_shape(r.chunks[0]), ((9, 6), (6, 6)))
 
@@ -54,6 +56,8 @@ def testQR(self):
 
         self.assertEqual(len(q.chunks), 1)
         self.assertEqual(len(r.chunks), 3)
+        self.assertEqual(q.nsplits, ((6,), (6,)))
+        self.assertEqual(r.nsplits, ((6,), (6, 6, 6)))
         self.assertEqual(calc_shape(q.chunks[0]), ((6, 6), (6, 6)))
         self.assertEqual(calc_shape(r.chunks[0]), ((6, 6), (6, 6)))
         self.assertEqual(calc_shape(r.chunks[1]), r.chunks[1].shape)
@@ -71,6 +75,8 @@ def testQR(self):
 
         self.assertEqual(len(q.chunks), 1)
         self.assertEqual(len(r.chunks), 2)
+        self.assertEqual(q.nsplits, ((6,), (6,)))
+        self.assertEqual(r.nsplits, ((6,), (6, 3)))
         self.assertEqual(calc_shape(q.chunks[0]), ((6, 6), (6, 6)))
         self.assertEqual(calc_shape(r.chunks[0]), ((6, 6), (6, 6)))
         self.assertEqual(calc_shape(r.chunks[1]), r.chunks[1].shape)
@@ -87,6 +93,8 @@ def testQR(self):
 
         self.assertEqual(len(q.chunks), 1)
         self.assertEqual(len(r.chunks), 1)
+        self.assertEqual(q.nsplits, ((9,), (6,)))
+        self.assertEqual(r.nsplits, ((6,), (6,)))
         self.assertEqual(calc_shape(q.chunks[0]), ((9, 6), (6, 6)))
         self.assertEqual(calc_shape(r.chunks[0]), ((9, 6), (6, 6)))