# Buggy code
def predict(self, data, **kw):
    session = kw.pop("session", None)
    run_kwargs = kw.pop("run_kwargs", dict())
    run = kw.pop("run", True)
    if kw:
        raise TypeError(
            f"predict got an unexpected keyword argument '{next(iter(kw))}'"
        )
    prob = predict(self.get_booster(), data, run=False)
    if prob.ndim > 1:
        prediction = mt.argmax(prob, axis=1)
    else:
        prediction = (prob > 0.5).astype(mt.int64)
    if run:
        prediction.execute(session=session, **run_kwargs)
    return prediction


# Traceback (ValueError)
In [1]: import mars.dataframe as md

In [2]: import numpy as np

In [3]: df = md.DataFrame({'a': np.random.randint(100, size=10), 'b':np.random.rand(10), 'label': np.random.randint(2, size=10)}, chunk_size=4)

In [4]: df = df[df.a > 0]

In [5]: df._shape = (10, 3)

In [6]: data = df[['a', 'b']]

In [7]: label = df['label']

In [8]: from mars.learn.contrib import xgboost as xgb

In [9]: xgb.MarsDMatrix(data=data, label=label).execute()
Out[9]:
Traceback (most recent call last):
File "/Users/qinxuye/miniconda3/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-13-43292b123748>", line 1, in <module>
xgb.MarsDMatrix(data=data, label=label).execute()
File "/Users/qinxuye/Workspace/mars/mars/learn/contrib/xgboost/dmatrix.py", line 327, in to_dmatrix
outs.execute(session=session, **(run_kwargs or dict()))
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 736, in execute
return super().execute(session=session, **kw)
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 379, in execute
return run()
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 374, in run
session.run(self, **kw)
File "/Users/qinxuye/Workspace/mars/mars/session.py", line 505, in run
result = self._sess.run(*tileables, **kw)
File "/Users/qinxuye/Workspace/mars/mars/session.py", line 111, in run
res = self._executor.execute_tileables(tileables, **kw)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/executor.py", line 866, in execute_tileables
tileables, tileable_graph=tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 201, in _tile
tds[0]._inplace_tile()
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 168, in _inplace_tile
return handler.inplace_tile(self)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 136, in inplace_tile
dispatched = self.dispatch(to_tile.op)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 119, in dispatch
tiled = op_cls.tile(op)
File "/Users/qinxuye/Workspace/mars/mars/learn/contrib/xgboost/dmatrix.py", line 255, in tile
return cls._tile_multi_output(op)
File "/Users/qinxuye/Workspace/mars/mars/learn/contrib/xgboost/dmatrix.py", line 139, in _tile_multi_output
label = label.rechunk({0: nsplit})._inplace_tile()
File "/Users/qinxuye/Workspace/mars/mars/tensor/rechunk/rechunk.py", line 97, in rechunk
chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
File "/Users/qinxuye/Workspace/mars/mars/tensor/rechunk/core.py", line 38, in get_nsplits
return decide_chunk_sizes(tileable.shape, chunk_size, itemsize)
File "/Users/qinxuye/Workspace/mars/mars/tensor/utils.py", line 551, in decide_chunk_sizes
return normalize_chunk_sizes(shape, tuple(chunk_size[j] for j in range(len(shape))))
File "/Users/qinxuye/Workspace/mars/mars/tensor/utils.py", line 66, in normalize_chunk_sizes
raise ValueError('chunks shape should be of the same length, '
ValueError: chunks shape should be of the same length, got shape: 10, chunks: (nan, nan, nan)

# Fixed code
def predict(self, data, **kw):
    session = kw.pop("session", None)
    run_kwargs = kw.pop("run_kwargs", dict())
    run = kw.pop("run", True)
    prob = predict(self.get_booster(), data, run=False, **kw)
    if prob.ndim > 1:
        prediction = mt.argmax(prob, axis=1)
    else:
        prediction = (prob > 0.5).astype(mt.int64)
    if run:
        prediction.execute(session=session, **run_kwargs)
    return prediction