# Buggy code
def tile(cls, op):
    """
    Use LU decomposition to compute inverse of matrix.
    Given a square matrix A:
    P, L, U = lu(A)
    b_eye is an identity matrix with the same shape as matrix A, then,
    (P * L * U) * A_inv = b_eye
    L * (U * A_inv) = P.T * b_eye
    use `solve_triangular` twice to compute the inverse of matrix A.
    """
    from .lu import lu
    from ..datasource import eye
    from ..base.transpose import TensorTranspose
    from .tensordot import tensordot
    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.single_tiles()

    p, l, u = lu(in_tensor)
    p.single_tiles()

    # transposed p equals to inverse of p
    p_transpose = TensorTranspose(
        dtype=p.dtype, sparse=p.op.sparse, axes=list(range(in_tensor.ndim))[::-1]
    ).new_tensor([p], p.shape)
    p_transpose.single_tiles()

    b = tensordot(p_transpose, b_eye, axes=((p_transpose.ndim - 1,), (b_eye.ndim - 2,)))
    b.single_tiles()

    # as `l` is a lower matrix, `lower=True` should be specified.
    uy = solve_triangular(l, b, lower=True)
    uy.single_tiles()

    a_inv = solve_triangular(u, uy)
    a_inv.single_tiles()
    return [a_inv]


# Traceback (KeyError)
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:

# Fixed code
def tile(cls, op):
    """
    Use LU decomposition to compute inverse of matrix.
    Given a square matrix A:
    P, L, U = lu(A)
    b_eye is an identity matrix with the same shape as matrix A, then,
    (P * L * U) * A_inv = b_eye
    L * (U * A_inv) = P.T * b_eye
    use `solve_triangular` twice to compute the inverse of matrix A.
    """
    from .lu import lu
    from ..datasource import eye
    from ..base.transpose import TensorTranspose
    from .tensordot import tensordot
    from .solve_triangular import solve_triangular

    in_tensor = op.input

    b_eye = eye(in_tensor.shape[0], chunk_size=in_tensor.nsplits)
    b_eye.single_tiles()

    p, l, u = lu(in_tensor)
    p.single_tiles()

    # transposed p equals to inverse of p
    p_transpose = TensorTranspose(
        dtype=p.dtype, sparse=p.op.sparse, axes=list(range(in_tensor.ndim))[::-1]
    ).new_tensor([p], p.shape)
    p_transpose.single_tiles()

    b = tensordot(p_transpose, b_eye, axes=((p_transpose.ndim - 1,), (b_eye.ndim - 2,)))
    b.single_tiles()

    # as `l` is a lower matrix, `lower=True` should be specified.
    uy = solve_triangular(l, b, lower=True)
    uy.single_tiles()

    a_inv = solve_triangular(u, uy)
    a_inv.single_tiles()
    return [a_inv]