# Buggy code
def predict(model, data, session=None, run_kwargs=None, run=True):
    from xgboost import Booster

    data = check_data(data)
    if not isinstance(model, Booster):
        raise TypeError(f"model has to be a xgboost.Booster, got {type(model)} instead")

    num_class = model.attr("num_class")
    if isinstance(data, TENSOR_TYPE):
        output_types = [OutputType.tensor]
    elif num_class is not None:
        output_types = [OutputType.dataframe]
    else:
        output_types = [OutputType.series]

    op = XGBPredict(data=data, model=model, gpu=data.op.gpu, output_types=output_types)
    result = op()
    if run:
        result.execute(session=session, **(run_kwargs or dict()))
    return result


# 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(
    model,
    data,
    output_margin=False,
    ntree_limit=None,
    validate_features=True,
    base_margin=None,
    session=None,
    run_kwargs=None,
    run=True,
):
    from xgboost import Booster

    data = check_data(data)
    if not isinstance(model, Booster):
        raise TypeError(f"model has to be a xgboost.Booster, got {type(model)} instead")

    num_class = model.attr("num_class")
    if isinstance(data, TENSOR_TYPE):
        output_types = [OutputType.tensor]
    elif num_class is not None:
        output_types = [OutputType.dataframe]
    else:
        output_types = [OutputType.series]

    kwargs = {
        "output_margin": output_margin,
        "ntree_limit": ntree_limit,
        "validate_features": validate_features,
        "base_margin": base_margin,
    }
    op = XGBPredict(
        data=data,
        model=model,
        kwargs=kwargs,
        gpu=data.op.gpu,
        output_types=output_types,
    )
    result = op()
    if run:
        result.execute(session=session, **(run_kwargs or dict()))
    return result