# Buggy code
def __call__(self):
    num_class = self._model.attr("num_class")
    if num_class is not None:
        num_class = int(num_class)
    if num_class is not None:
        shape = (self._data.shape[0], num_class)
    else:
        shape = (self._data.shape[0],)
    if self.output_types[0] == OutputType.tensor:
        # tensor
        return self.new_tileable(
            [self._data],
            shape=shape,
            dtype=np.dtype(np.float32),
            order=TensorOrder.C_ORDER,
        )
    elif self.output_types[0] == OutputType.dataframe:
        # dataframe
        dtypes = pd.DataFrame(np.random.rand(0, num_class), dtype=np.float32).dtypes
        return self.new_tileable(
            [self._data],
            shape=shape,
            dtypes=dtypes,
            columns_value=parse_index(dtypes.index),
            index_value=self._data.index_value,
        )
    else:
        # series
        return self.new_tileable(
            [self._data],
            shape=shape,
            index_value=self._data.index_value,
            name="predictions",
            dtype=np.dtype(np.float32),
        )


# 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 __call__(self):
    num_class = self._model.attr("num_class")
    if num_class is not None:
        num_class = int(num_class)
    if num_class is not None:
        shape = (self._data.shape[0], num_class)
    else:
        shape = (self._data.shape[0],)
    inputs = [self._data]
    if self.output_types[0] == OutputType.tensor:
        # tensor
        return self.new_tileable(
            inputs, shape=shape, dtype=np.dtype(np.float32), order=TensorOrder.C_ORDER
        )
    elif self.output_types[0] == OutputType.dataframe:
        # dataframe
        dtypes = pd.DataFrame(np.random.rand(0, num_class), dtype=np.float32).dtypes
        return self.new_tileable(
            inputs,
            shape=shape,
            dtypes=dtypes,
            columns_value=parse_index(dtypes.index),
            index_value=self._data.index_value,
        )
    else:
        # series
        return self.new_tileable(
            inputs,
            shape=shape,
            index_value=self._data.index_value,
            name="predictions",
            dtype=np.dtype(np.float32),
        )