# Buggy code
def tile(cls, op: "LGBMAlign"):
    inputs = [
        d for d in [op.data, op.label, op.sample_weight, op.init_score] if d is not None
    ]
    data = op.data

    ctx = get_context()
    if ctx.running_mode != RunningMode.distributed:
        outputs = [
            inp.rechunk(tuple((s,) for s in inp.shape))._inplace_tile()
            for inp in inputs
        ]
    else:
        if len(data.nsplits[1]) != 1:
            data = data.rechunk({1: data.shape[1]})._inplace_tile()
        outputs = [data]
        for inp in inputs[1:]:
            if inp is not None:
                outputs.append(inp.rechunk((data.nsplits[0],))._inplace_tile())

    kws = []
    for o in outputs:
        kw = o.params.copy()
        kw.update(dict(chunks=o.chunks, nsplits=o.nsplits))
        kws.append(kw)

    new_op = op.copy().reset_key()
    tileables = new_op.new_tileables(inputs, kws=kws)

    return tileables


# Traceback (ValueError)
Traceback (most recent call last):
File "/home/admin/work/_public-mars-0.4.2.zip/mars/scheduler/graph.py", line 371, in _execute_graph
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 343, in _wrapped
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 483, in inner
File "/home/admin/work/_public-mars-0.4.2.zip/mars/scheduler/graph.py", line 603, in prepare_graph
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 389, in _wrapped
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 483, in inner
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 342, in build
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 389, in _wrapped
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 483, in inner
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 255, in build
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 294, in inner
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 235, in build
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 330, in _tile
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 195, in _tile
File "/home/admin/work/_public-mars-0.4.2.zip/mars/scheduler/graph.py", line 586, in on_tile
File "/home/admin/work/_public-mars-0.4.2.zip/mars/context.py", line 68, in h
File "/home/admin/work/_public-mars-0.4.2.zip/mars/utils.py", line 389, in _wrapped
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tiles.py", line 115, in dispatch
File "/home/admin/work/_public-mars-0.4.2.zip/mars/learn/contrib/lightgbm/align.py", line 87, in tile
File "/home/admin/work/_public-mars-0.4.2.zip/mars/dataframe/base/rechunk.py", line 90, in rechunk
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tensor/rechunk/core.py", line 38, in get_nsplits
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tensor/utils.py", line 549, in decide_chunk_sizes
File "/home/admin/work/_public-mars-0.4.2.zip/mars/tensor/utils.py", line 67, in normalize_chunk_sizes
ValueError: chunks shape should be of the same length, got shape: nan, chunks: (nan,)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/admin/work/_public-mars-0.4.2.zip/mars/promise.py", line 100, in _wrapped
File "/home/admin/work/_public-mars-0.4.2.zip/mars/worker/calc.py", line 271, in <lambda>
File "/home/admin/work/_public-mars-0.4.2.zip/mars/worker/calc.py", line 244, in _start_calc
File "/home/admin/work/_public-mars-0.4.2.zip/mars/worker/calc.py", line 173, in _calc_results
File "src/gevent/event.py", line 383, in gevent._gevent_cevent.AsyncResult.result
File "src/gevent/event.py", line 305, in gevent._gevent_cevent.AsyncResult.get
File "src/gevent/event.py", line 335, in gevent._gevent_cevent.AsyncResult.get
File "src/gevent/event.py", line 323, in gevent._gevent_cevent.AsyncResult.get
File "src/gevent/event.py", line 303, in gevent._gevent_cevent.AsyncResult._raise_exception
File "/opt/conda/lib/python3.7/site-packages/gevent/_compat.py", line 65, in reraise
File "/opt/conda/lib/python3.7/site-packages/gevent/threadpool.py", line 142, in __run_task
File "mars/actors/pool/gevent_pool.pyx", line 127, in mars.actors.pool.gevent_pool.GeventThreadPool._wrap_watch.inner
result = fn(*args, **kwargs)
File "/home/admin/work/_public-mars-0.4.2.zip/mars/executor.py", line 690, in execute_graph
File "/home/admin/work/_public-mars-0.4.2.zip/mars/executor.py", line 571, in execute
File "/opt/conda/lib/python3.7/concurrent/futures/_base.py", line 435, in result
File "/opt/conda/lib/python3.7/concurrent/futures/_base.py", line 384, in __get_result
File "/opt/conda/lib/python3.7/concurrent/futures/thread.py", line 57, in run
File "/home/admin/work/_public-mars-0.4.2.zip/mars/executor.py", line 443, in _execute_operand
File "/home/admin/work/_public-mars-0.4.2.zip/mars/executor.py", line 641, in handle
File "/home/admin/work/_public-pyodps-0.9.3.zip/odps/mars_extension/core.py", line 187, in wrapper
File "/home/admin/work/_public-mars-0.4.2.zip/mars/remote/core.py", line 200, in execute
File "E:/MyDocuments/PycharmProjects/borrowloan/temp_query/test_mars.py", line 31, in light_gbm
lg_reg = lgb.LGBMRegressor()
File "/home/admin/work/_public-mars-0.4.2.zip/mars/learn/contrib/lightgbm/regressor.py", line 28, in fit
File "/home/admin/work/_public-mars-0.4.2.zip/mars/learn/contrib/lightgbm/train.py", line 322, in train
File "/home/admin/work/_public-mars-0.4.2.zip/mars/core.py", line 651, in execute
File "/home/admin/work/_public-mars-0.4.2.zip/mars/core.py", line 370, in execute
File "/home/admin/work/_public-mars-0.4.2.zip/mars/session.py", line 428, in run
File "/home/admin/work/_public-mars-0.4.2.zip/mars/session.py", line 303, in run
mars.errors.ExecutionFailed: '\'\\\'\\\\\\\'"\\\\\\\\\\\\\\\'Graph execution failed.\\\\\\\\\\\\\\\'"\\\\\\\'\\\'\''

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "E:/MyDocuments/PycharmProjects/borrowloan/temp_query/test_mars.py", line 38, in <module>
odps.run_mars_job(light_gbm, args=(tb_name,), worker_num=2, worker_cpu=2, worker_mem=8, mars_image='extended',
File "F:\Anaconda3\lib\site-packages\odps\mars_extension\core.py", line 151, in run_mars_job
r.execute()
File "F:\Anaconda3\lib\site-packages\mars\core.py", line 370, in execute
session.run(self, **kw)
File "F:\Anaconda3\lib\site-packages\mars\session.py", line 428, in run
result = self._sess.run(*tileables, **kw)
File "F:\Anaconda3\lib\site-packages\mars\web\session.py", line 187, in run
if self._check_response_finished(graph_url, timeout_val):
File "F:\Anaconda3\lib\site-packages\mars\web\session.py", line 146, in _check_response_finished
raise ExecutionFailed('Graph execution failed.') from exc
mars.errors.ExecutionFailed: 'Graph execution failed.'

# Fixed code
def tile(cls, op: "LGBMAlign"):
    inputs = [
        d for d in [op.data, op.label, op.sample_weight, op.init_score] if d is not None
    ]
    data = op.data

    # check inputs to make sure no unknown chunk shape exists
    check_chunks_unknown_shape(inputs, TilesError)

    ctx = get_context()
    if ctx.running_mode != RunningMode.distributed:
        outputs = [
            inp.rechunk(tuple((s,) for s in inp.shape))._inplace_tile()
            for inp in inputs
        ]
    else:
        if len(data.nsplits[1]) != 1:
            data = data.rechunk({1: data.shape[1]})._inplace_tile()
        outputs = [data]
        for inp in inputs[1:]:
            if inp is not None:
                outputs.append(inp.rechunk((data.nsplits[0],))._inplace_tile())

    kws = []
    for o in outputs:
        kw = o.params.copy()
        kw.update(dict(chunks=o.chunks, nsplits=o.nsplits))
        kws.append(kw)

    new_op = op.copy().reset_key()
    tileables = new_op.new_tileables(inputs, kws=kws)

    return tileables