# Buggy code
def rechunk(
    a, chunk_size, threshold=None, chunk_size_limit=None, reassign_worker=False
):
    if isinstance(a, DATAFRAME_TYPE):
        itemsize = max(getattr(dt, "itemsize", 8) for dt in a.dtypes)
    else:
        itemsize = a.dtype.itemsize
    chunk_size = get_nsplits(a, chunk_size, itemsize)
    if chunk_size == a.nsplits:
        return a

    op = DataFrameRechunk(
        chunk_size, threshold, chunk_size_limit, reassign_worker=reassign_worker
    )
    return op(a)


# Traceback (ValueError)
Traceback (most recent call last):
File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/graph.py", line 410, in _execute_graph
self.prepare_graph(compose=compose)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 377, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/graph.py", line 648, in prepare_graph
self._target_tileable_datas + fetch_tileables, tileable_graph)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 203, in _tile
tds = on_tile(tileable_data.op.outputs, tds)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/graph.py", line 630, in on_tile
return self.context.wraps(handler.dispatch)(first.op)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/context.py", line 72, in h
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 119, in dispatch
tiled = op_cls.tile(op)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/learn/contrib/xgboost/dmatrix.py", line 255, in tile
return cls._tile_multi_output(op)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/learn/contrib/xgboost/dmatrix.py", line 139, in _tile_multi_output
label = label.rechunk({0: nsplit})._inplace_tile()
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tensor/rechunk/rechunk.py", line 81, in rechunk
chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
File "/home/admin/work/_public-mars-0.6.2.zip/mars/tensor/rechunk/core.py", line 38, in get_nsplits
return decide_chunk_sizes(tileable.shape, chunk_size, itemsize)
File "/home/admin/work/_public-mars-0.6.2.zip/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 "/home/admin/work/_public-mars-0.6.2.zip/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: 1000000, chunks: (nan,)

# Fixed code
def rechunk(
    a, chunk_size, threshold=None, chunk_size_limit=None, reassign_worker=False
):
    if not any(pd.isna(s) for s in a.shape):
        # do client check only when no unknown shape,
        # real nsplits will be recalculated inside `tile`
        chunk_size = _get_chunk_size(a, chunk_size)
        if chunk_size == a.nsplits:
            return a

    op = DataFrameRechunk(
        chunk_size=chunk_size,
        threshold=threshold,
        chunk_size_limit=chunk_size_limit,
        reassign_worker=reassign_worker,
    )
    return op(a)