# Buggy code
def _wrapped(*args, **kwargs):
    _kernel_mode.eager = False
    return_value = func(*args, **kwargs)
    _kernel_mode.eager = None
    return return_value


# Traceback (KeyError)
In [1]: import mars.tensor as mt

In [2]: from mars.config import options

In [3]: a = mt.array(
...:     [[0.1, 0.2, 0.3],
...:      [0.3, 0.4, 0.2]]
...:     )

In [4]: (a[:, mt.newaxis, :] - a[mt.newaxis, ...]).execute()
Out[4]:
array([[[ 0. ,  0. ,  0. ],
[-0.2, -0.2,  0.1]],

[[ 0.2,  0.2, -0.1],
[ 0. ,  0. ,  0. ]]])

In [5]: options.eager_mode = True

In [6]:  mt.sqrt(mt.sum((a[:, mt.newaxis, :] - a[mt.newaxis, ...]) ** 2, axis=-1))
---------------------------------------------------------------------------
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.arithmetic.sqrt.TensorSqrt'>

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-6-de975d9f2cf8> in <module>
----> 1 mt.sqrt(mt.sum((a[:, mt.newaxis, :] - a[mt.newaxis, ...]) ** 2, axis=-1))

~/Workspace/mars/mars/tensor/expressions/utils.py in h(*tensors, **kw)
157                 kw['dtype'] = dtype
158
--> 159             ret = func(*tensors, **kw)
160             if ret is NotImplemented:
161                 reverse_func = getattr(inspect.getmodule(func), 'r{0}'.format(func.__name__), None) \

~/Workspace/mars/mars/tensor/expressions/arithmetic/sqrt.py in sqrt(x, out, where, **kwargs)
78     """
79     op = TensorSqrt(**kwargs)
---> 80     return op(x, out=out, where=where)
81

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in __call__(self, x, out, where)
348
349     def __call__(self, x, out=None, where=None):
--> 350         return self._call(x, out=out, where=where)
351
352

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in _call(self, x, out, where)
331         shape = x.shape
332
--> 333         t = self.new_tensor([x, out, where], shape)
334
335         if out is None:

~/Workspace/mars/mars/tensor/expressions/core.py in new_tensor(self, inputs, shape, dtype, **kw)
54             raise TypeError('cannot new tensor with more than 1 outputs')
55
---> 56         return self.new_tensors(inputs, shape, dtype=dtype, **kw)[0]
57
58     def calc_shape(self, *inputs_shape):

~/Workspace/mars/mars/tensor/expressions/core.py in new_tensors(self, inputs, shape, dtype, chunks, nsplits, output_limit, kws, **kw)
48                     output_limit=None, kws=None, **kw):
49         return self.new_entities(inputs, shape, chunks=chunks, nsplits=nsplits,
---> 50                                  output_limit=output_limit, kws=kws, dtype=dtype, **kw)
51
52     def new_tensor(self, inputs, shape, dtype=None, **kw):

~/Workspace/mars/mars/core.py in new_entities(self, inputs, shape, **kwargs)
578         entities = self._new_entities(inputs, shape, **kwargs)
579         if is_eager_mode():
--> 580             ExecutableTuple(entities).execute(fetch=False)
581         return entities
582

~/Workspace/mars/mars/core.py in execute(self, session, **kw)
594         if session is None:
595             session = Session.default_or_local()
--> 596         return session.run(*self, **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/tensor/core.py in tiles(self)
223
224     def tiles(self):
--> 225         return handler.tiles(self)
226
227     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/arithmetic/core.py in tile(cls, op)
43         for out_index in itertools.product(*(map(range, out_chunk_shape))):
44             in_chunks = [t.cix[get_index(out_index[-t.ndim:], t)] if t.ndim != 0 else t.chunks[0]
---> 45                          for t in inputs]
46             chunk_op = op.copy().reset_key()
47             chunk_shape = broadcast_shape(*(c.shape for c in in_chunks))

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in <listcomp>(.0)
43         for out_index in itertools.product(*(map(range, out_chunk_shape))):
44             in_chunks = [t.cix[get_index(out_index[-t.ndim:], t)] if t.ndim != 0 else t.chunks[0]
---> 45                          for t in inputs]
46             chunk_op = op.copy().reset_key()
47             chunk_shape = broadcast_shape(*(c.shape for c in in_chunks))

~/Workspace/mars/mars/core.py in __getitem__(self, item)
457                 if len(item) != self._tilesable.ndim:
458                     raise ValueError('Cannot get tensor chunk by %s, expect length %d' % (
--> 459                         item, self._tilesable.ndim))
460
461                 s = self._tilesable.chunk_shape

ValueError: Cannot get tensor chunk by (0, 0), expect length 3

In [7]: a[:, mt.newaxis, :] - a[mt.newaxis, ...]
Out[7]: Tensor <op=TensorSubtract, shape=(2, 2, 3), key=fe6279614688ba864daa99368caf12af>

# Fixed code
def _wrapped(*args, **kwargs):
    try:
        _kernel_mode.eager = False
        return func(*args, **kwargs)
    finally:
        _kernel_mode.eager = None