Repository: mars-project/mars

Issue:
Remove limitation in cholesky decomposition
Cholesky decomposition's tile logic is actually similar to lu decomposition, in #272 , the limitation that all chunks must be square was removed, so should Cholesky.


Discussion:

Resolving patch:
diff --git a/mars/tensor/expressions/linalg/cholesky.py b/mars/tensor/expressions/linalg/cholesky.py
--- a/mars/tensor/expressions/linalg/cholesky.py
+++ b/mars/tensor/expressions/linalg/cholesky.py
@@ -56,12 +56,10 @@ def tile(cls, op):
 
         tensor = op.outputs[0]
         in_tensor = op.input
-        if len(set(in_tensor.nsplits[0] + in_tensor.nsplits[1])) != 1:
-            nsplit = min(in_tensor.nsplits[0] + in_tensor.nsplits[1])
-            # input's chunks must be all square
-            in_tensor = in_tensor.rechunk([nsplit for _ in range(in_tensor.ndim)]).single_tiles()
-            if len(set(in_tensor.nsplits[0] + in_tensor.nsplits[1])) != 1:
-                raise LinAlgError('All chunks must be a square matrix to perform cholesky decomposition.')
+        if in_tensor.nsplits[0] != in_tensor.nsplits[1]:
+            # all chunks on diagonal should be square
+            nsplits = in_tensor.nsplits[0]
+            in_tensor = in_tensor.rechunk([nsplits, nsplits]).single_tiles()
 
         lower_chunks, upper_chunks = {}, {}
         for i in range(in_tensor.chunk_shape[0]):
@@ -70,7 +68,7 @@ def tile(cls, op):
                     lower_chunk = TensorZeros(dtype=tensor.dtype).new_chunk(
                         None, (in_tensor.nsplits[0][i], in_tensor.nsplits[1][j]), index=(i, j))
                     upper_chunk = TensorZeros(dtype=tensor.dtype).new_chunk(
-                        None, (in_tensor.nsplits[0][i], in_tensor.nsplits[1][j]), index=(j, i))
+                        None, (in_tensor.nsplits[1][j], in_tensor.nsplits[0][i]), index=(j, i))
                     lower_chunks[lower_chunk.index] = lower_chunk
                     upper_chunks[upper_chunk.index] = upper_chunk
                 elif i == j:


Test patch:
diff --git a/mars/tensor/execution/tests/test_linalg_execute.py b/mars/tensor/execution/tests/test_linalg_execute.py
--- a/mars/tensor/execution/tests/test_linalg_execute.py
+++ b/mars/tensor/execution/tests/test_linalg_execute.py
@@ -122,40 +122,68 @@ def testSVDExecution(self):
         np.testing.assert_array_almost_equal(res, expected)
 
     def testCholeskyExecution(self):
-        data = np.array([[1, -2j], [2j, 5]])
+        data = np.random.randint(1, 10, (10, 10))
+        symmetric_data = data.dot(data.T)
 
-        a = tensor(data, chunk_size=1)
+        a = tensor(symmetric_data, chunk_size=5)
+
+        U = cholesky(a)
+        t = U.T.dot(U)
+
+        res_u = self.executor.execute_tensor(U, concat=True)[0]
+        np.testing.assert_allclose(np.triu(res_u), res_u)
+
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        self.assertTrue(np.allclose(res, symmetric_data))
 
         L = cholesky(a, lower=True)
-        t = L.dot(L.T.conj())
+        U = cholesky(a)
+        t = L.dot(U)
 
         res = self.executor.execute_tensor(t, concat=True)[0]
-        self.assertTrue(np.allclose(res, data))
+        self.assertTrue(np.allclose(res, symmetric_data))
+
+        a = tensor(symmetric_data, chunk_size=2)
 
         L = cholesky(a, lower=True)
         U = cholesky(a)
         t = L.dot(U)
 
         res = self.executor.execute_tensor(t, concat=True)[0]
-        self.assertTrue(np.allclose(res, data))
+        np.testing.assert_allclose(res, symmetric_data)
 
-        a = tensor(data, chunk_size=2)
+        a = tensor(symmetric_data, chunk_size=(1, 2))
 
         L = cholesky(a, lower=True)
         U = cholesky(a)
         t = L.dot(U)
 
         res = self.executor.execute_tensor(t, concat=True)[0]
-        np.testing.assert_allclose(res, data)
+        np.testing.assert_allclose(res, symmetric_data)
 
-        a = tensor(data, chunk_size=(1, 2))
+        a = tensor(symmetric_data, chunk_size=4)
 
         L = cholesky(a, lower=True)
         U = cholesky(a)
         t = L.dot(U)
 
+        res_u = self.executor.execute_tensor(U, concat=True)[0]
+        np.testing.assert_allclose(np.triu(res_u), res_u)
+
         res = self.executor.execute_tensor(t, concat=True)[0]
-        np.testing.assert_allclose(res, data)
+        np.testing.assert_allclose(res, symmetric_data)
+
+        a = tensor(symmetric_data, chunk_size=3)
+
+        L = cholesky(a, lower=True)
+        U = cholesky(a)
+        t = L.dot(U)
+
+        res_u = self.executor.execute_tensor(U, concat=True)[0]
+        np.testing.assert_allclose(np.triu(res_u), res_u)
+
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        np.testing.assert_allclose(res, symmetric_data)
 
     def testLUExecution(self):
         np.random.seed(1)