# Buggy code
def tile(cls, op):
    tensor = op.tensor
    pk = op.pk
    out = op.outputs[0]
    index_path = op.index_path
    ctx = get_context()

    # check index_path for distributed
    if getattr(ctx, "running_mode", None) == RunningMode.distributed:
        if index_path is not None:
            fs = get_fs(index_path, op.storage_options)
            if isinstance(fs, LocalFileSystem):
                raise ValueError(
                    "`index_path` cannot be local file dir for distributed index building"
                )

    # make sure all inputs have known chunk sizes
    check_chunks_unknown_shape(op.inputs, TilesError)

    nsplit = decide_unify_split(tensor.nsplits[0], pk.nsplits[0])
    if op.topk is not None:
        nsplit = cls._get_atleast_topk_nsplit(nsplit, op.topk)

    if tensor.chunk_shape[1] > 1:
        tensor = tensor.rechunk({0: nsplit, 1: tensor.shape[1]})._inplace_tile()
    else:
        tensor = tensor.rechunk({0: nsplit})._inplace_tile()
    pk = pk.rechunk({0: nsplit})._inplace_tile()

    out_chunks = []
    for chunk, pk_col_chunk in zip(tensor.chunks, pk.chunks):
        chunk_op = op.copy().reset_key()
        chunk_op._stage = OperandStage.map
        out_chunk = chunk_op.new_chunk([chunk, pk_col_chunk], index=pk_col_chunk.index)
        out_chunks.append(out_chunk)

    logger.warning(f"index chunks count: {len(out_chunks)} ")

    params = out.params
    params["chunks"] = out_chunks
    params["nsplits"] = ((1,) * len(out_chunks),)
    new_op = op.copy()
    return new_op.new_tileables(op.inputs, kws=[params])


# Traceback (ValueError)
Traceback (most recent call last):
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/utils.py", line 365, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/scheduler/graph.py", line 382, in execute_graph
self._execute_graph(compose=compose)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/scheduler/graph.py", line 410, in _execute_graph
self.prepare_graph(compose=compose)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/utils.py", line 365, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/utils.py", line 439, in _inner
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/scheduler/graph.py", line 648, in prepare_graph
self._target_tileable_datas + fetch_tileables, tileable_graph)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/utils.py", line 439, in _inner
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/utils.py", line 439, in _inner
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 203, in _tile
tds = on_tile(tileable_data.op.outputs, tds)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/scheduler/graph.py", line 630, in on_tile
return self.context.wraps(handler.dispatch)(first.op)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/context.py", line 72, in h
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/utils.py", line 439, in _inner
return func(*args, **kwargs)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tiles.py", line 119, in dispatch
tiled = op_cls.tile(op)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/dataframe/merge/concat.py", line 182, in tile
return cls._tile_dataframe(op)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/dataframe/merge/concat.py", line 94, in _tile_dataframe
inputs = [item.rechunk(normalized_nsplits)._inplace_tile() for item in inputs]
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/dataframe/merge/concat.py", line 94, in <listcomp>
inputs = [item.rechunk(normalized_nsplits)._inplace_tile() for item in inputs]
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/dataframe/base/rechunk.py", line 96, in rechunk
chunk_size = get_nsplits(a, chunk_size, itemsize)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.zip/mars/tensor/rechunk/core.py", line 38, in get_nsplits
return decide_chunk_sizes(tileable.shape, chunk_size, itemsize)
File "/home/admin/work/tip_dev-pymars-0.6.0a3.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/tip_dev-pymars-0.6.0a3.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: 100000, chunks: (nan, nan, nan, nan, nan, nan)

# Fixed code
def tile(cls, op):
    tensor = op.tensor
    pk = op.pk
    out = op.outputs[0]
    index_path = op.index_path
    ctx = get_context()
    fs = None
    if index_path is not None:
        fs = get_fs(index_path, op.storage_options)

    # check index_path for distributed
    if getattr(ctx, "running_mode", None) == RunningMode.distributed:
        if index_path is not None:
            if isinstance(fs, LocalFileSystem):
                raise ValueError(
                    "`index_path` cannot be local file dir "
                    "for distributed index building"
                )

    if index_path is not None:
        # check if the index path is empty
        try:
            files = [f for f in fs.ls(index_path) if "proxima_" in f]
            if files:
                raise ValueError(
                    f"Directory {index_path} contains built proxima index, "
                    f"clean them to perform new index building"
                )
        except FileNotFoundError:
            # if not exist, create directory
            fs.mkdir(index_path)

    # make sure all inputs have known chunk sizes
    check_chunks_unknown_shape(op.inputs, TilesError)

    nsplit = decide_unify_split(tensor.nsplits[0], pk.nsplits[0])
    if op.topk is not None:
        nsplit = cls._get_atleast_topk_nsplit(nsplit, op.topk)

    if tensor.chunk_shape[1] > 1:
        tensor = tensor.rechunk({0: nsplit, 1: tensor.shape[1]})._inplace_tile()
    else:
        tensor = tensor.rechunk({0: nsplit})._inplace_tile()
    pk = pk.rechunk({0: nsplit})._inplace_tile()

    out_chunks = []
    for chunk, pk_col_chunk in zip(tensor.chunks, pk.chunks):
        chunk_op = op.copy().reset_key()
        chunk_op._stage = OperandStage.map
        out_chunk = chunk_op.new_chunk([chunk, pk_col_chunk], index=pk_col_chunk.index)
        out_chunks.append(out_chunk)

    logger.warning(f"index chunks count: {len(out_chunks)} ")

    params = out.params
    params["chunks"] = out_chunks
    params["nsplits"] = ((1,) * len(out_chunks),)
    new_op = op.copy()
    return new_op.new_tileables(op.inputs, kws=[params])