Repository: mars-project/mars

Issue:
[BUG] Chunk's shape size is incorrect in SVD
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

The shape of SVD chunk is incorrect in some cases.

**To Reproduce**
```
In [1]: import mars.tensor as mt                                                                                                    

In [2]: a = mt.random.rand(10, 50, chunk_size=20)                                                                                   

In [3]: u, s, v = mt.linalg.svd(a)                                                                                                  

In [4]: u.tiles()                                                                                                                   
Out[4]: Tensor <op=TensorSVD, shape=(10, 10), key=dc4fe7c18a97d7943c196de086119bad>

In [5]: print(u.shape, u.chunks[0].shape)                                                                                           
(10, 10) (10, 50)

```


Discussion:

Resolving patch:
diff --git a/mars/tensor/linalg/core.py b/mars/tensor/linalg/core.py
--- a/mars/tensor/linalg/core.py
+++ b/mars/tensor/linalg/core.py
@@ -19,6 +19,7 @@
 from ...compat import izip
 from ..core import TensorOrder
 from ..utils import decide_chunk_sizes
+from .utils import calc_svd_shapes
 
 
 class SFQR(object):
@@ -166,9 +167,7 @@ def tile(cls, op):
             U_shape, s_shape, V_shape = U.shape, s.shape, V.shape
 
             svd_op = TensorSVD()
-            u_shape = stage2_r_chunk.shape
-            s_shape = (stage2_r_chunk.shape[1],)
-            v_shape = (stage2_r_chunk.shape[1],) * 2
+            u_shape, s_shape, v_shape = calc_svd_shapes(stage2_r_chunk)
             stage2_usv_chunks = svd_op.new_chunks([stage2_r_chunk],
                                                   kws=[{'side': 'U', 'dtype': U_dtype,
                                                         'index': stage2_r_chunk.index,
diff --git a/mars/tensor/linalg/svd.py b/mars/tensor/linalg/svd.py
--- a/mars/tensor/linalg/svd.py
+++ b/mars/tensor/linalg/svd.py
@@ -25,6 +25,7 @@
 from ..operands import TensorHasInput, TensorOperandMixin
 from ..core import TensorOrder
 from .core import TSQR
+from .utils import calc_svd_shapes
 
 
 class TensorSVD(TensorHasInput, TensorOperandMixin):
@@ -63,15 +64,7 @@ def __call__(self, a):
 
         # if a's shape is (6, 18), U's shape is (6, 6), s's shape is (6,), V's shape is (6, 18)
         # if a's shape is (18, 6), U's shape is (18, 6), s's shape is (6,), V's shape is (6, 6)
-        x, y = a.shape
-        if x > y:
-            U_shape = (x, y)
-            s_shape = (y, )
-            V_shape = (y, y)
-        else:
-            U_shape = (x, x)
-            s_shape = (x, )
-            V_shape = (x, y)
+        U_shape, s_shape, V_shape = calc_svd_shapes(a)
         U, s, V = self.new_tensors([a],
                                    order=TensorOrder.C_ORDER,
                                    kws=[
diff --git a/mars/tensor/linalg/utils.py b/mars/tensor/linalg/utils.py
--- a/mars/tensor/linalg/utils.py
+++ b/mars/tensor/linalg/utils.py
@@ -15,6 +15,22 @@
 import numpy as np
 
 
+def calc_svd_shapes(a):
+    """
+    Calculate output shapes of singular value decomposition.
+    Follow the behavior of `numpy`:
+    if a's shape is (6, 18), U's shape is (6, 6), s's shape is (6,), V's shape is (6, 18)
+    if a's shape is (18, 6), U's shape is (18, 6), s's shape is (6,), V's shape is (6, 6)
+    :param a: input tensor
+    :return: (U.shape, s.shape, V.shape)
+    """
+    x, y = a.shape
+    if x > y:
+        return (x, y), (y,), (y, y)
+    else:
+        return (x, x), (x,), (x, y)
+
+
 def svd_flip(u, v, u_based_decision=True):
     """
     Sign correction to ensure deterministic output from SVD.


Test patch:
diff --git a/mars/tensor/linalg/tests/test_linalg.py b/mars/tensor/linalg/tests/test_linalg.py
--- a/mars/tensor/linalg/tests/test_linalg.py
+++ b/mars/tensor/linalg/tests/test_linalg.py
@@ -107,8 +107,11 @@ def testSVD(self):
 
         U.tiles()
         self.assertEqual(len(U.chunks), 3)
+        self.assertEqual(U.chunks[0].shape, (3, 6))
         self.assertEqual(len(s.chunks), 1)
+        self.assertEqual(s.chunks[0].shape, (6,))
         self.assertEqual(len(V.chunks), 1)
+        self.assertEqual(V.chunks[0].shape, (6, 6))
 
         self.assertEqual(s.ndim, 1)
         self.assertEqual(len(s.chunks[0].index), 1)
@@ -122,12 +125,30 @@ def testSVD(self):
 
         U.tiles()
         self.assertEqual(len(U.chunks), 1)
+        self.assertEqual(U.chunks[0].shape, (9, 6))
         self.assertEqual(len(s.chunks), 1)
+        self.assertEqual(s.chunks[0].shape, (6,))
         self.assertEqual(len(V.chunks), 1)
+        self.assertEqual(V.chunks[0].shape, (6, 6))
 
         self.assertEqual(s.ndim, 1)
         self.assertEqual(len(s.chunks[0].index), 1)
 
+        a = mt.random.rand(6, 20, chunk_size=10)
+        U, s, V = mt.linalg.svd(a)
+
+        self.assertEqual(U.shape, (6, 6))
+        self.assertEqual(s.shape, (6,))
+        self.assertEqual(V.shape, (6, 20))
+
+        U.tiles()
+        self.assertEqual(len(U.chunks), 1)
+        self.assertEqual(U.chunks[0].shape, (6, 6))
+        self.assertEqual(len(s.chunks), 1)
+        self.assertEqual(s.chunks[0].shape, (6,))
+        self.assertEqual(len(V.chunks), 1)
+        self.assertEqual(V.chunks[0].shape, (6, 20))
+
         a = mt.random.rand(6, 9, chunk_size=(6, 9))
         U, s, V = mt.linalg.svd(a)