Repository: mars-project/mars

Issue:
Remove the limitation in lu decomposition
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**
`mt.linalg.lu` is an important part in Mars linear algebra module which is used by `mt.linalg.inv` and `mt.linalg.solve`, however it has a limitation that all of input matrix's chunks should be square. Suppose the length of matrix is a prime number, it can't be split into chunks that all of them are square matrix, and it will no be executed in parallel by Mars.

``` Python
In [1]: import mars.tensor as mt                                                                                                         

In [2]: a = mt.random.rand(8, 8, chunk_size=4)                                                                                           

In [3]: mt.linalg.inv(a).execute()                                                                                                       
Out[3]: 
array([[ 1.03846764e+00,  2.07539634e+00, -4.77775050e-01,
        -2.44298670e-01,  1.77448669e-01, -1.27463956e+00,
        -3.05228117e+00,  1.11515300e+00],
       [-1.43125778e+00,  1.75146909e-01, -5.78229471e-01,
         1.54459643e-01,  1.38989770e+00,  6.46797260e-02,
         3.32975869e-01,  4.55595401e-02],
       [ 7.52315493e-01, -1.79781805e+00, -9.11656770e-01,
         5.97371314e-01, -4.75625170e-01,  1.38827913e+00,
         2.50281780e+00, -1.16138635e+00],
       [-4.95582404e-01, -9.89574604e-02,  2.95866331e-01,
        -3.18312861e-01,  2.58706766e-01,  1.37111292e+00,
        -2.94282309e-01, -1.16138810e-01],
       [ 1.01840500e+00, -1.39759384e+00,  5.67016260e-01,
         9.04352464e-01, -7.92237836e-01, -3.17215562e-01,
         3.94887006e-01, -2.34236135e-01],
       [-1.17750293e+00,  3.32287224e+00, -1.10667783e+00,
        -4.06912889e-03,  1.60029178e+00,  4.68348180e-01,
        -2.90109865e+00, -3.73910491e-01],
       [-1.09861743e-02, -7.77764078e-01, -1.02812075e-01,
        -3.08869352e-01,  5.12033373e-01,  2.15181685e-03,
         2.12120064e+00, -9.82795797e-01],
       [ 3.29387228e-01,  6.53558413e-02,  2.20095703e+00,
        -3.66117635e-01, -1.81340471e+00, -1.84828247e+00,
        -6.61805623e-01,  2.08285783e+00]])

In [4]: a = mt.random.rand(8, 8, chunk_size=5)                                                                                           

In [5]: mt.linalg.inv(a).execute()                                                                                                       
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/Documents/mars_dev/mars/mars/tiles.py in _dispatch(self, op)
    110         try:
--> 111             handler = self._handlers[op_cls]
    112             return handler(op)

KeyError: <class 'mars.tensor.expressions.linalg.inv.TensorInv'>

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
~/Documents/mars_dev/mars/mars/tiles.py in _dispatch(self, op)
    110         try:
--> 111             handler = self._handlers[op_cls]
    112             return handler(op)

KeyError: <class 'mars.tensor.expressions.linalg.lu.TensorLU'>

During handling of the above exception, another exception occurred:

LinAlgError                               Traceback (most recent call last)
<ipython-input-5-0fba0aa06805> in <module>
----> 1 mt.linalg.inv(a).execute()

~/Documents/mars_dev/mars/mars/tensor/core.py in execute(self, session, **kw)
    378 
    379     def execute(self, session=None, **kw):
--> 380         return self._data.execute(session, **kw)
    381 
    382 

~/Documents/mars_dev/mars/mars/tensor/core.py in execute(self, session, **kw)
    184         if session is None:
    185             session = Session.default_or_local()
--> 186         return session.run(self, **kw)
    187 
    188     def fetch(self, session=None, **kw):

~/Documents/mars_dev/mars/mars/session.py in run(self, *tensors, **kw)
    107 
    108         tensors = tuple(mt.tensor(t) for t in tensors)
--> 109         result = self._sess.run(*tensors, **kw)
    110 
    111         for t in tensors:

~/Documents/mars_dev/mars/mars/session.py in run(self, *tensors, **kw)
     49         if 'n_parallel' not in kw:
     50             kw['n_parallel'] = cpu_count()
---> 51         res = self._executor.execute_tensors(tensors, **kw)
     52         return res
     53 

~/Documents/mars_dev/mars/mars/utils.py in _wrapped(*args, **kwargs)
    352     def _wrapped(*args, **kwargs):
    353         _kernel_mode.eager = False
--> 354         return_value = func(*args, **kwargs)
    355         _kernel_mode.eager = None
    356         return return_value

~/Documents/mars_dev/mars/mars/executor.py in execute_tensors(self, tensors, fetch, n_parallel, n_thread, print_progress, mock, sparse_mock_percent)
    451         concat_keys = []
    452         for tensor in tensors:
--> 453             tensor.tiles()
    454             chunk_keys = [c.key for c in tensor.chunks]
    455             result_keys.extend(chunk_keys)

~/Documents/mars_dev/mars/mars/core.py in tiles(self)
    390 
    391     def tiles(self):
--> 392         return handler.tiles(self)
    393 
    394     def single_tiles(self):

~/Documents/mars_dev/mars/mars/tiles.py in tiles(self, tiles_obj)
    172             if not preds or accessible:
    173                 if node.is_coarse() and node.op:
--> 174                     tiled = self._dispatch(node.op)
    175                     self._assign_to([t.data for t in tiled], node.op.outputs)
    176                 visited.add(node)

~/Documents/mars_dev/mars/mars/utils.py in _wrapped(*args, **kwargs)
    352     def _wrapped(*args, **kwargs):
    353         _kernel_mode.eager = False
--> 354         return_value = func(*args, **kwargs)
    355         _kernel_mode.eager = None
    356         return return_value

~/Documents/mars_dev/mars/mars/tiles.py in _dispatch(self, op)
    114             if hasattr(op_cls, 'tile'):
    115                 # has tile implementation
--> 116                 return op_cls.tile(op)
    117             for op_clz in self._handlers.keys():
    118                 if issubclass(op_cls, op_clz):

~/Documents/mars_dev/mars/mars/tensor/expressions/linalg/inv.py in tile(cls, op)
     57 
     58         p, l, u = lu(in_tensor)
---> 59         p.single_tiles()
     60 
     61         # transposed p equals to inverse of p

~/Documents/mars_dev/mars/mars/tensor/core.py in single_tiles(self)
    226 
    227     def single_tiles(self):
--> 228         return handler.single_tiles(self)
    229 
    230     @property

~/Documents/mars_dev/mars/mars/tiles.py in single_tiles(self, to_tiles)
    127     def single_tiles(self, to_tiles):
    128         if to_tiles.is_coarse() and to_tiles.op:
--> 129             dispatched = self._dispatch(to_tiles.op)
    130             self._assign_to([d.data for d in dispatched], to_tiles.op.outputs)
    131 

~/Documents/mars_dev/mars/mars/utils.py in _wrapped(*args, **kwargs)
    352     def _wrapped(*args, **kwargs):
    353         _kernel_mode.eager = False
--> 354         return_value = func(*args, **kwargs)
    355         _kernel_mode.eager = None
    356         return return_value

~/Documents/mars_dev/mars/mars/tiles.py in _dispatch(self, op)
    114             if hasattr(op_cls, 'tile'):
    115                 # has tile implementation
--> 116                 return op_cls.tile(op)
    117             for op_clz in self._handlers.keys():
    118                 if issubclass(op_cls, op_clz):

~/Documents/mars_dev/mars/mars/tensor/expressions/linalg/lu.py in tile(cls, op)
     71             in_tensor = in_tensor.rechunk([nsplit for _ in range(in_tensor.ndim)]).single_tiles()
     72             if len(set(in_tensor.nsplits[0] + in_tensor.nsplits[1])) != 1:
---> 73                 raise LinAlgError('All chunks must be a square matrix to perform LU decomposition.')
     74 
     75         p_chunks, p_invert_chunks, lower_chunks, l_permuted_chunks, upper_chunks = {}, {}, {}, {}, {}

LinAlgError: All chunks must be a square matrix to perform LU decomposition.
```
**Describe the solution you'd like**
Actually, the tile logic in `mt.linalg.lu` is work for the situation mentioned above, it can be executed even not all chunks are square.



Discussion:

Resolving patch:
diff --git a/mars/tensor/expressions/linalg/lu.py b/mars/tensor/expressions/linalg/lu.py
--- a/mars/tensor/expressions/linalg/lu.py
+++ b/mars/tensor/expressions/linalg/lu.py
@@ -65,12 +65,11 @@ def tile(cls, op):
 
         P, L, U = op.outputs
         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 LU 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()
 
         p_chunks, p_invert_chunks, lower_chunks, l_permuted_chunks, upper_chunks = {}, {}, {}, {}, {}
         for i in range(in_tensor.chunk_shape[0]):


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
@@ -161,34 +161,60 @@ def testLUExecution(self):
         np.random.seed(1)
 
         data = np.random.randint(1, 10, (6, 6))
-        a = tensor(data, chunk_size=2)
 
+        a = tensor(data)
         P, L, U = lu(a)
-        t = P.dot(L).dot(U)
 
+        # check lower and upper triangular matrix
+        result_l = self.executor.execute_tensor(L, concat=True)[0]
+        result_u = self.executor.execute_tensor(U, concat=True)[0]
+
+        np.testing.assert_allclose(np.tril(result_l), result_l)
+        np.testing.assert_allclose(np.triu(result_u), result_u)
+
+        t = P.dot(L).dot(U)
         res = self.executor.execute_tensor(t, concat=True)[0]
         np.testing.assert_allclose(res, data)
 
-        a = tensor(data, chunk_size=1)
-
+        a = tensor(data, chunk_size=2)
         P, L, U = lu(a)
-        t = P.dot(L).dot(U)
 
+        # check lower and upper triangular matrix
+        result_l = self.executor.execute_tensor(L, concat=True)[0]
+        result_u = self.executor.execute_tensor(U, concat=True)[0]
+
+        np.testing.assert_allclose(np.tril(result_l), result_l)
+        np.testing.assert_allclose(np.triu(result_u), result_u)
+
+        t = P.dot(L).dot(U)
         res = self.executor.execute_tensor(t, concat=True)[0]
         np.testing.assert_allclose(res, data)
 
-        a = tensor(data, chunk_size=(1, 2))
-
+        a = tensor(data, chunk_size=(2, 3))
         P, L, U = lu(a)
-        t = P.dot(L).dot(U)
 
+        # check lower and upper triangular matrix
+        result_l = self.executor.execute_tensor(L, concat=True)[0]
+        result_u = self.executor.execute_tensor(U, concat=True)[0]
+
+        np.testing.assert_allclose(np.tril(result_l), result_l)
+        np.testing.assert_allclose(np.triu(result_u), result_u)
+
+        t = P.dot(L).dot(U)
         res = self.executor.execute_tensor(t, concat=True)[0]
         np.testing.assert_allclose(res, data)
 
-        a = tensor(data, chunk_size=3)
+        a = tensor(data, chunk_size=4)
         P, L, U = lu(a)
-        t = P.dot(L).dot(U)
 
+        # check lower and upper triangular matrix
+        result_l = self.executor.execute_tensor(L, concat=True)[0]
+        result_u = self.executor.execute_tensor(U, concat=True)[0]
+
+        np.testing.assert_allclose(np.tril(result_l), result_l)
+        np.testing.assert_allclose(np.triu(result_u), result_u)
+
+        t = P.dot(L).dot(U)
         res = self.executor.execute_tensor(t, concat=True)[0]
         np.testing.assert_allclose(res, data)
 
@@ -199,6 +225,22 @@ def testLUExecution(self):
                                [0, 6, 0, 8, 7, 3],
                                [7, 0, 6, 1, 7, 0],
                                [0, 0, 0, 7, 0, 8]])
+
+        a = tensor(data)
+        P, L, U = lu(a)
+        result_l = self.executor.execute_tensor(L, concat=True)[0]
+        result_u = self.executor.execute_tensor(U, concat=True)[0]
+
+        # check lower and upper triangular matrix
+        np.testing.assert_allclose(np.tril(result_l), result_l)
+        np.testing.assert_allclose(np.triu(result_u), result_u)
+        self.assertIsInstance(result_l, SparseNDArray)
+        self.assertIsInstance(result_u, SparseNDArray)
+
+        t = P.dot(L).dot(U)
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        np.testing.assert_array_almost_equal(data.A, res)
+
         a = tensor(data, chunk_size=2)
         P, L, U = lu(a)
         result_l = self.executor.execute_tensor(L, concat=True)[0]
@@ -214,7 +256,7 @@ def testLUExecution(self):
         res = self.executor.execute_tensor(t, concat=True)[0]
         np.testing.assert_array_almost_equal(data.A, res)
 
-        a = tensor(data, chunk_size=1)
+        a = tensor(data, chunk_size=(2, 3))
         P, L, U = lu(a)
         result_l = self.executor.execute_tensor(L, concat=True)[0]
         result_u = self.executor.execute_tensor(U, concat=True)[0]
@@ -229,7 +271,7 @@ def testLUExecution(self):
         res = self.executor.execute_tensor(t, concat=True)[0]
         np.testing.assert_array_almost_equal(data.A, res)
 
-        a = tensor(data, chunk_size=6)
+        a = tensor(data, chunk_size=4)
         P, L, U = lu(a)
         result_l = self.executor.execute_tensor(L, concat=True)[0]
         result_u = self.executor.execute_tensor(U, concat=True)[0]
@@ -384,6 +426,29 @@ def testSolve(self):
         res = self.executor.execute_tensor(A.dot(x), concat=True)[0]
         np.testing.assert_allclose(res, data2)
 
+        # test for not all chunks are square in matrix A
+        data2 = np.random.randint(1, 10, (20,))
+
+        A = tensor(data1, chunk_size=6)
+        b = tensor(data2, chunk_size=6)
+
+        x = solve(A, b)
+
+        res = self.executor.execute_tensor(x, concat=True)[0]
+        np.testing.assert_allclose(res, scipy.linalg.solve(data1, data2))
+        res = self.executor.execute_tensor(A.dot(x), concat=True)[0]
+        np.testing.assert_allclose(res, data2)
+
+        A = tensor(data1, chunk_size=(7, 6))
+        b = tensor(data2, chunk_size=6)
+
+        x = solve(A, b)
+
+        res = self.executor.execute_tensor(x, concat=True)[0]
+        np.testing.assert_allclose(res, scipy.linalg.solve(data1, data2))
+        res = self.executor.execute_tensor(A.dot(x), concat=True)[0]
+        np.testing.assert_allclose(res, data2)
+
         # test sparse
         data1 = sps.csr_matrix(np.random.randint(1, 10, (20, 20)))
         data2 = np.random.randint(1, 10, (20, ))
@@ -419,6 +484,17 @@ def testSolve(self):
         self.assertIsInstance(res, SparseNDArray)
         np.testing.assert_allclose(res, data2)
 
+        # test for not all chunks are square in matrix A
+        data2 = np.random.randint(1, 10, (20,))
+
+        A = tensor(data1, chunk_size=6)
+        b = tensor(data2, chunk_size=6)
+
+        x = solve(A, b)
+
+        res = self.executor.execute_tensor(A.dot(x), concat=True)[0]
+        np.testing.assert_allclose(res, data2)
+
     def testSolveSymPos(self):
         import scipy.linalg
         np.random.seed(1)
@@ -467,6 +543,15 @@ def testInv(self):
         res = self.executor.execute_tensor(B.dot(inv_B), concat=True)[0]
         self.assertTrue(np.allclose(res, np.eye(data.shape[0], dtype=float)))
 
+        # test for not all chunks are square in matrix A
+        A = tensor(data, chunk_size=8)
+        inv_A = inv(A)
+
+        res = self.executor.execute_tensor(inv_A, concat=True)[0]
+        self.assertTrue(np.allclose(res, scipy.linalg.inv(data)))
+        res = self.executor.execute_tensor(A.dot(inv_A), concat=True)[0]
+        self.assertTrue(np.allclose(res, np.eye(data.shape[0], dtype=float)))
+
         # test sparse
         data = np.random.randint(1, 10, (20, 20))
         sp_data = sps.csr_matrix(data)
@@ -480,6 +565,16 @@ def testInv(self):
         res = self.executor.execute_tensor(A.dot(inv_A), concat=True)[0]
         self.assertTrue(np.allclose(res, np.eye(data.shape[0], dtype=float)))
 
+        # test for not all chunks are square in matrix A
+        A = tensor(sp_data, chunk_size=8)
+        inv_A = inv(A)
+
+        res = self.executor.execute_tensor(inv_A, concat=True)[0]
+        self.assertIsInstance(res, SparseNDArray)
+        self.assertTrue(np.allclose(res, scipy.linalg.inv(data)))
+        res = self.executor.execute_tensor(A.dot(inv_A), concat=True)[0]
+        self.assertTrue(np.allclose(res, np.eye(data.shape[0], dtype=float)))
+
     def testNormExecution(self):
         d = np.arange(9) - 4
         d2 = d.reshape(3, 3)
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
@@ -205,6 +205,30 @@ def testLU(self):
         self.assertEqual(calc_shape(l.chunks[0]), l.chunks[0].shape)
         self.assertEqual(calc_shape(u.chunks[0]), u.chunks[0].shape)
 
+        a = mt.random.randint(1, 10, (6, 6), chunk_size=(3, 2))
+        p, l, u = mt.linalg.lu(a)
+        l.tiles()
+
+        self.assertEqual(l.shape, (6, 6))
+        self.assertEqual(u.shape, (6, 6))
+        self.assertEqual(p.shape, (6, 6))
+
+        self.assertEqual(p.nsplits, ((3, 3), (3, 3)))
+        self.assertEqual(l.nsplits, ((3, 3), (3, 3)))
+        self.assertEqual(u.nsplits, ((3, 3), (3, 3)))
+
+        a = mt.random.randint(1, 10, (7, 7), chunk_size=4)
+        p, l, u = mt.linalg.lu(a)
+        l.tiles()
+
+        self.assertEqual(l.shape, (7, 7))
+        self.assertEqual(u.shape, (7, 7))
+        self.assertEqual(p.shape, (7, 7))
+
+        self.assertEqual(p.nsplits, ((4, 3), (4, 3)))
+        self.assertEqual(l.nsplits, ((4, 3), (4, 3)))
+        self.assertEqual(u.nsplits, ((4, 3), (4, 3)))
+
         # test sparse
         data = sps.csr_matrix([[2, 0, 0, 0, 5, 2],
                                [0, 6, 1, 0, 0, 6],
@@ -243,6 +267,13 @@ def testSolve(self):
         self.assertEqual(calc_shape(x), x.shape)
         self.assertEqual(calc_shape(x.chunks[0]), x.chunks[0].shape)
 
+        a = mt.random.randint(1, 10, (20, 20), chunk_size=12)
+        b = mt.random.randint(1, 10, (20, 3))
+        x = mt.linalg.solve(a, b).tiles()
+
+        self.assertEqual(x.shape, (20, 3))
+        self.assertEqual(x.nsplits, ((12, 8), (3, )))
+
         # test sparse
         a = sps.csr_matrix(np.random.randint(1, 10, (20, 20)))
         b = mt.random.randint(1, 10, (20, ), chunk_size=3)
@@ -253,6 +284,13 @@ def testSolve(self):
         self.assertTrue(x.op.sparse)
         self.assertTrue(x.chunks[0].op.sparse)
 
+        a = mt.tensor(a, chunk_size=7)
+        b = mt.random.randint(1, 10, (20,))
+        x = mt.linalg.solve(a, b).tiles()
+
+        self.assertEqual(x.shape, (20,))
+        self.assertEqual(x.nsplits, ((7, 7, 6),))
+
         x = mt.linalg.solve(a, b, sparse=False).tiles()
         self.assertFalse(x.op.sparse)
         self.assertFalse(x.chunks[0].op.sparse)
@@ -265,6 +303,12 @@ def testInv(self):
         self.assertEqual(calc_shape(a_inv), a_inv.shape)
         self.assertEqual(calc_shape(a_inv.chunks[0]), a_inv.chunks[0].shape)
 
+        a = mt.random.randint(1, 10, (20, 20), chunk_size=11)
+        a_inv = mt.linalg.inv(a).tiles()
+
+        self.assertEqual(a_inv.shape, (20, 20))
+        self.assertEqual(a_inv.nsplits, ((11, 9), (11, 9)))
+
         b = a.T.dot(a)
         b_inv = mt.linalg.inv(b).tiles()
         self.assertEqual(b_inv.shape, (20, 20))