# Buggy code
def build_index(
    tensor,
    pk,
    dimension=None,
    index_path=None,
    need_shuffle=False,
    distance_metric="SquaredEuclidean",
    index_builder="SsgBuilder",
    index_builder_params=None,
    index_converter=None,
    index_converter_params=None,
    topk=None,
    storage_options=None,
    run=True,
    session=None,
    run_kwargs=None,
):
    tensor = validate_tensor(tensor)

    if dimension is None:
        dimension = tensor.shape[1]
    if index_builder_params is None:
        index_builder_params = {}
    if index_converter_params is None:
        index_converter_params = {}

    if need_shuffle:
        tensor = mt.random.permutation(tensor)

    if not isinstance(pk, (Base, Entity)):
        pk = mt.tensor(pk)

    op = ProximaBuilder(
        tensor=tensor,
        pk=pk,
        distance_metric=distance_metric,
        index_path=index_path,
        dimension=dimension,
        index_builder=index_builder,
        index_builder_params=index_builder_params,
        index_converter=index_converter,
        index_converter_params=index_converter_params,
        topk=topk,
        storage_options=storage_options,
    )
    result = op(tensor, pk)
    if run:
        return result.execute(session=session, **(run_kwargs or dict()))
    else:
        return result


# 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 build_index(
    tensor,
    pk,
    dimension=None,
    index_path=None,
    need_shuffle=False,
    distance_metric="SquaredEuclidean",
    index_builder="SsgBuilder",
    index_builder_params=None,
    index_converter=None,
    index_converter_params=None,
    topk=None,
    storage_options=None,
    run=True,
    session=None,
    run_kwargs=None,
):
    tensor = validate_tensor(tensor)
    if tensor.dtype not in available_numpy_dtypes:
        raise ValueError(
            f"Dtype to build index should be one of {available_numpy_dtypes}, "
            f"got {tensor.dtype}"
        )

    if dimension is None:
        dimension = tensor.shape[1]
    if index_builder_params is None:
        index_builder_params = {}
    if index_converter_params is None:
        index_converter_params = {}

    if need_shuffle:
        tensor = mt.random.permutation(tensor)

    if not isinstance(pk, (Base, Entity)):
        pk = mt.tensor(pk)

    op = ProximaBuilder(
        tensor=tensor,
        pk=pk,
        distance_metric=distance_metric,
        index_path=index_path,
        dimension=dimension,
        index_builder=index_builder,
        index_builder_params=index_builder_params,
        index_converter=index_converter,
        index_converter_params=index_converter_params,
        topk=topk,
        storage_options=storage_options,
    )
    result = op(tensor, pk)
    if run:
        return result.execute(session=session, **(run_kwargs or dict()))
    else:
        return result