# Buggy code
def tile(cls, op: "ProximaSearcher"):
    tensor = op.tensor
    index = op.index
    topk = op.topk
    outs = op.outputs

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

    if tensor.chunk_shape[1] > 1:
        tensor = tensor.rechunk({1: tensor.shape[1]})._inplace_tile()

    logger.warning(f"query chunks count: {len(tensor.chunks)} ")

    if hasattr(index, "op"):
        built_indexes = index.chunks
    else:
        # index path
        fs: FileSystem = get_fs(index, op.storage_options)
        built_indexes = [
            f for f in fs.ls(index) if f.rsplit("/", 1)[-1].startswith("proxima-")
        ]

    if hasattr(index, "op"):
        ctx = get_context()
        index_chunks_workers = [
            m.workers[0] if m.workers else None
            for m in ctx.get_chunk_metas([c.key for c in index.chunks])
        ]
    else:
        index_chunks_workers = [None] * len(built_indexes)

    out_chunks = [], []
    for tensor_chunk in tensor.chunks:
        pk_chunks, distance_chunks = [], []
        for j, chunk_index, worker in zip(
            itertools.count(), built_indexes, index_chunks_workers
        ):
            chunk_op = op.copy().reset_key()
            chunk_op._stage = OperandStage.map
            if hasattr(chunk_index, "op"):
                chunk_op._expect_worker = worker
            chunk_op._index = chunk_index
            chunk_kws = [
                {
                    "index": (tensor_chunk.index[0], j),
                    "dtype": outs[0].dtype,
                    "shape": (tensor_chunk.shape[0], topk),
                    "order": TensorOrder.C_ORDER,
                },
                {
                    "index": (tensor_chunk.index[0], j),
                    "dtype": outs[1].dtype,
                    "shape": (tensor_chunk.shape[0], topk),
                    "order": TensorOrder.C_ORDER,
                },
            ]
            chunk_inputs = [tensor_chunk]
            if hasattr(chunk_index, "op"):
                chunk_inputs.append(chunk_index)
            pk_chunk, distance_chunk = chunk_op.new_chunks(chunk_inputs, kws=chunk_kws)
            pk_chunks.append(pk_chunk)
            distance_chunks.append(distance_chunk)

        if len(pk_chunks) == 1:
            out_chunks[0].append(pk_chunks[0])
            out_chunks[1].append(distance_chunks[0])
            continue

        shape = (tensor_chunk.shape[0], topk * len(pk_chunks))
        pk_merge_op = TensorConcatenate(axis=1)
        pk_merge_chunk = pk_merge_op.new_chunk(
            pk_chunks,
            index=(pk_chunks[0].index[0], 0),
            shape=shape,
            dtype=pk_chunks[0].dtype,
            order=pk_chunks[0].order,
        )
        distance_merge_op = TensorConcatenate(axis=1)
        distance_merge_chunk = distance_merge_op.new_chunk(
            distance_chunks,
            index=(distance_chunks[0].index[0], 0),
            shape=shape,
            dtype=distance_chunks[0].dtype,
            order=distance_chunks[0].order,
        )

        agg_op = ProximaSearcher(
            stage=OperandStage.agg, topk=op.topk, distance_metric=op.distance_metric
        )
        agg_chunk_kws = [
            {
                "index": pk_merge_chunk.index,
                "dtype": outs[0].dtype,
                "shape": (tensor_chunk.shape[0], topk),
                "order": outs[0].order,
            },
            {
                "index": pk_merge_chunk.index,
                "dtype": outs[1].dtype,
                "shape": (tensor_chunk.shape[0], topk),
                "order": outs[1].order,
            },
        ]
        pk_result_chunk, distance_result_chunk = agg_op.new_chunks(
            [pk_merge_chunk, distance_merge_chunk], kws=agg_chunk_kws
        )
        out_chunks[0].append(pk_result_chunk)
        out_chunks[1].append(distance_result_chunk)

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

    kws = []
    pk_params = outs[0].params
    pk_params["chunks"] = out_chunks[0]
    pk_params["nsplits"] = (tensor.nsplits[0], (topk,))
    kws.append(pk_params)
    distance_params = outs[1].params
    distance_params["chunks"] = out_chunks[1]
    distance_params["nsplits"] = (tensor.nsplits[0], (topk,))
    kws.append(distance_params)
    new_op = op.copy()
    return new_op.new_tileables(op.inputs, kws=kws)


# 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: "ProximaSearcher"):
    tensor = op.tensor
    index = op.index
    topk = op.topk
    outs = op.outputs

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

    if tensor.chunk_shape[1] > 1:
        tensor = tensor.rechunk({1: tensor.shape[1]})._inplace_tile()

    logger.warning(f"query chunks count: {len(tensor.chunks)} ")

    if hasattr(index, "op"):
        built_indexes = index.chunks
    else:
        # index path
        fs: FileSystem = get_fs(index, op.storage_options)
        built_indexes = [
            f for f in fs.ls(index) if f.rsplit("/", 1)[-1].startswith("proxima_")
        ]

    if hasattr(index, "op"):
        ctx = get_context()
        index_chunks_workers = [
            m.workers[0] if m.workers else None
            for m in ctx.get_chunk_metas([c.key for c in index.chunks])
        ]
    else:
        index_chunks_workers = [None] * len(built_indexes)

    out_chunks = [], []
    for tensor_chunk in tensor.chunks:
        pk_chunks, distance_chunks = [], []
        for j, chunk_index, worker in zip(
            itertools.count(), built_indexes, index_chunks_workers
        ):
            chunk_op = op.copy().reset_key()
            chunk_op._stage = OperandStage.map
            if hasattr(chunk_index, "op"):
                chunk_op._expect_worker = worker
            chunk_op._index = chunk_index
            chunk_kws = [
                {
                    "index": (tensor_chunk.index[0], j),
                    "dtype": outs[0].dtype,
                    "shape": (tensor_chunk.shape[0], topk),
                    "order": TensorOrder.C_ORDER,
                },
                {
                    "index": (tensor_chunk.index[0], j),
                    "dtype": outs[1].dtype,
                    "shape": (tensor_chunk.shape[0], topk),
                    "order": TensorOrder.C_ORDER,
                },
            ]
            chunk_inputs = [tensor_chunk]
            if hasattr(chunk_index, "op"):
                chunk_inputs.append(chunk_index)
            pk_chunk, distance_chunk = chunk_op.new_chunks(chunk_inputs, kws=chunk_kws)
            pk_chunks.append(pk_chunk)
            distance_chunks.append(distance_chunk)

        if len(pk_chunks) == 1:
            out_chunks[0].append(pk_chunks[0])
            out_chunks[1].append(distance_chunks[0])
            continue

        shape = (tensor_chunk.shape[0], topk * len(pk_chunks))
        pk_merge_op = TensorConcatenate(axis=1)
        pk_merge_chunk = pk_merge_op.new_chunk(
            pk_chunks,
            index=(pk_chunks[0].index[0], 0),
            shape=shape,
            dtype=pk_chunks[0].dtype,
            order=pk_chunks[0].order,
        )
        distance_merge_op = TensorConcatenate(axis=1)
        distance_merge_chunk = distance_merge_op.new_chunk(
            distance_chunks,
            index=(distance_chunks[0].index[0], 0),
            shape=shape,
            dtype=distance_chunks[0].dtype,
            order=distance_chunks[0].order,
        )

        agg_op = ProximaSearcher(
            stage=OperandStage.agg, topk=op.topk, distance_metric=op.distance_metric
        )
        agg_chunk_kws = [
            {
                "index": pk_merge_chunk.index,
                "dtype": outs[0].dtype,
                "shape": (tensor_chunk.shape[0], topk),
                "order": outs[0].order,
            },
            {
                "index": pk_merge_chunk.index,
                "dtype": outs[1].dtype,
                "shape": (tensor_chunk.shape[0], topk),
                "order": outs[1].order,
            },
        ]
        pk_result_chunk, distance_result_chunk = agg_op.new_chunks(
            [pk_merge_chunk, distance_merge_chunk], kws=agg_chunk_kws
        )
        out_chunks[0].append(pk_result_chunk)
        out_chunks[1].append(distance_result_chunk)

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

    kws = []
    pk_params = outs[0].params
    pk_params["chunks"] = out_chunks[0]
    pk_params["nsplits"] = (tensor.nsplits[0], (topk,))
    kws.append(pk_params)
    distance_params = outs[1].params
    distance_params["chunks"] = out_chunks[1]
    distance_params["nsplits"] = (tensor.nsplits[0], (topk,))
    kws.append(distance_params)
    new_op = op.copy()
    return new_op.new_tileables(op.inputs, kws=kws)