fix chunk index of tensor of singular values calculated by svd decomp…

…osition (#58)
mars-project · Dec 20, 2018 · 599819f · 599819f
1 parent 2ff2cfe
commit 599819f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 8 deletions.
diff --git a/mars/tensor/expressions/linalg/core.py b/mars/tensor/expressions/linalg/core.py
@@ -116,10 +116,12 @@ def tile(cls, op):
             s_shape = (stage2_r_chunk.shape[1],)
             v_shape = (stage2_r_chunk.shape[1],) * 2
             stage2_usv_chunks = svd_op.new_chunks([stage2_r_chunk], [u_shape, s_shape, v_shape],
-                                                  index=stage2_r_chunk.index,
-                                                  kws=[{'side': 'U', 'dtype': U_dtype},
-                                                       {'side': 's', 'dtype': s_dtype},
-                                                       {'side': 'V', 'dtype': V_dtype}])
+                                                  kws=[{'side': 'U', 'dtype': U_dtype,
+                                                        'index': stage2_r_chunk.index},
+                                                       {'side': 's', 'dtype': s_dtype,
+                                                        'index': stage2_r_chunk.index[1:]},
+                                                       {'side': 'V', 'dtype': V_dtype,
+                                                        'index': stage2_r_chunk.index}])
             stage2_u_chunk, stage2_s_chunk, stage2_v_chunk = stage2_usv_chunks
 
             # stage 4, U = Q @ u

diff --git a/mars/tensor/expressions/linalg/svd.py b/mars/tensor/expressions/linalg/svd.py
@@ -68,11 +68,13 @@ def tile(cls, op):
             in_chunk = in_tensor.chunks[0]
             chunk_op = op.copy().reset_key()
             svd_chunks = chunk_op.new_chunks([in_chunk], (U_shape, s_shape, V_shape),
-                                             index=in_chunk.index,
                                              kws=[
-                                                 {'side': 'U', 'dtype': U_dtype},
-                                                 {'side': 's', 'dtype': s_dtype},
-                                                 {'side': 'V', 'dtype': V_dtype}
+                                                 {'side': 'U', 'dtype': U_dtype,
+                                                  'index': in_chunk.index},
+                                                 {'side': 's', 'dtype': s_dtype,
+                                                  'index': in_chunk.index[1:]},
+                                                 {'side': 'V', 'dtype': V_dtype,
+                                                  'index': in_chunk.index}
                                              ])
             U_chunk, s_chunk, V_chunk = svd_chunks
 

diff --git a/mars/tensor/expressions/tests/test_linalg.py b/mars/tensor/expressions/tests/test_linalg.py
@@ -63,6 +63,24 @@ def testSVD(self):
         self.assertEqual(len(s.chunks), 1)
         self.assertEqual(len(V.chunks), 1)
 
+        self.assertEqual(s.ndim, 1)
+        self.assertEqual(len(s.chunks[0].index), 1)
+
+        a = mt.random.rand(9, 6, chunks=(9, 6))
+        U, s, V = mt.linalg.svd(a)
+
+        self.assertEqual(U.shape, (9, 6))
+        self.assertEqual(s.shape, (6,))
+        self.assertEqual(V.shape, (6, 6))
+
+        U.tiles()
+        self.assertEqual(len(U.chunks), 1)
+        self.assertEqual(len(s.chunks), 1)
+        self.assertEqual(len(V.chunks), 1)
+
+        self.assertEqual(s.ndim, 1)
+        self.assertEqual(len(s.chunks[0].index), 1)
+
         rs = mt.random.RandomState(1)
         a = rs.rand(9, 6, chunks=(3, 6))
         U, s, V = mt.linalg.svd(a)