Repository: mars-project/mars

Issue:
[BUG] Cannot execute mt.linalg.inv given a sparse input tensor
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Given an input tensor which is created from `scipy.sparse.csr_matrix`, cannot perform execution of `mt.linalg.inv`.

**To Reproduce**

```
In [32]: a = mt.tensor(sps.csr_matrix([[0, 0], [1, 0]]))                            

In [33]: b = mt.linalg.inv(a)                                                       

In [34]: b.execute()                                                                
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/Workspace/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)
~/Workspace/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.datasource.eye.TensorEye'>

During handling of the above exception, another exception occurred:

AssertionError                            Traceback (most recent call last)
<ipython-input-34-be8c8c42b9b4> in <module>
----> 1 b.execute()

~/Workspace/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 

~/Workspace/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):

~/Workspace/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:

~/Workspace/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 

~/Workspace/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

~/Workspace/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)

~/Workspace/mars/mars/core.py in tiles(self)
    390 
    391     def tiles(self):
--> 392         return handler.tiles(self)
    393 
    394     def single_tiles(self):

~/Workspace/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)

~/Workspace/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

~/Workspace/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):

~/Workspace/mars/mars/tensor/expressions/linalg/inv.py in tile(cls, op)
     53 
     54         b_eye = eye(in_tensor.shape[0], chunk_size=in_tensor.chunks[0].shape)
---> 55         b_eye.single_tiles()
     56 
     57         p, l, u = lu(in_tensor)

~/Workspace/mars/mars/tensor/core.py in single_tiles(self)
    226 
    227     def single_tiles(self):
--> 228         return handler.single_tiles(self)
    229 
    230     @property

~/Workspace/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 

~/Workspace/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

~/Workspace/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):

~/Workspace/mars/mars/tensor/expressions/datasource/eye.py in tile(cls, op)
     51     @classmethod
     52     def tile(cls, op):
---> 53         return TensorDiagBase.tile(op)
     54 
     55 

~/Workspace/mars/mars/tensor/expressions/datasource/diag.py in tile(cls, op)
     61         # op can be TensorDiag or TensorEye
     62         k = op.k
---> 63         nsplits = op._get_nsplits(op)
     64 
     65         fx = lambda x, y: x - y + k

~/Workspace/mars/mars/tensor/expressions/datasource/eye.py in _get_nsplits(cls, op)
     42         tensor = op.outputs[0]
     43         chunk_size = tensor.params.raw_chunk_size or options.tensor.chunk_size
---> 44         return decide_chunk_sizes(tensor.shape, chunk_size, tensor.dtype.itemsize)
     45 
     46     @classmethod

~/Workspace/mars/mars/tensor/expressions/utils.py in decide_chunk_sizes(shape, chunk_size, itemsize)
    434         raise ValueError("chunks have more dimensions than input tensor")
    435     if nleft == 0:
--> 436         return normalize_chunk_sizes(shape, tuple(chunk_size[j] for j in range(len(shape))))
    437 
    438     max_chunk_size = options.tensor.chunk_store_limit

~/Workspace/mars/mars/tensor/expressions/utils.py in normalize_chunk_sizes(shape, chunk_size)
     57             chunk_sizes.append(chunk)
     58         else:
---> 59             assert isinstance(chunk, six.integer_types)
     60 
     61             sizes = tuple(chunk for _ in range(int(size / chunk))) + \

AssertionError: 
```


Discussion:

Resolving patch:
diff --git a/mars/tensor/expressions/linalg/inv.py b/mars/tensor/expressions/linalg/inv.py
--- a/mars/tensor/expressions/linalg/inv.py
+++ b/mars/tensor/expressions/linalg/inv.py
@@ -51,7 +51,7 @@ def tile(cls, op):
         from .solve_triangular import solve_triangular
         in_tensor = op.input
 
-        b_eye = eye(in_tensor.shape[0], chunk_size=in_tensor.chunks[0].shape)
+        b_eye = eye(in_tensor.shape[0], chunk_size=in_tensor.nsplits)
         b_eye.single_tiles()
 
         p, l, u = lu(in_tensor)


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
@@ -433,6 +433,14 @@ def testInv(self):
 
         data = np.random.randint(1, 10, (20, 20))
 
+        A = tensor(data)
+        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)))
+
         A = tensor(data, chunk_size=5)
         inv_A = inv(A)