Repository: dask/dask
File: dask/array/_shuffle.py (Python)

Code before:
    >>> y.compute()
    array([[ 7,  6,  3,  5,  2,  4,  1,  8],
           [15, 14, 11, 13, 10, 12,  9, 16]])
    """
    if np.isnan(x.shape).any():
        raise ValueError(
            f"Shuffling only allowed with known chunk sizes. {unknown_chunk_message}"
        )
    assert isinstance(axis, int), "axis must be an integer"
    _validate_indexer(x.chunks, indexer, axis)

    x = _resize_other_dimensions(x, max(map(len, indexer)), axis, chunks)

    token = tokenize(x, indexer, axis)
    out_name = f"shuffle-{token}"

    chunks, layer = _shuffle(x.chunks, indexer, axis, x.name, out_name, token)
    if len(layer) == 0:
        return Array(x.dask, x.name, x.chunks, meta=x)

    graph = HighLevelGraph.from_collections(out_name, layer, dependencies=[x])

    return Array(graph, out_name, chunks, meta=x)


def _resize_other_dimensions(x: Array, longest_group, axis, chunks):
    assert chunks == "auto", "Only auto is supported for now"
    chunksize_tolerance = config.get("array.shuffle.chunksize-tolerance")

    if longest_group <= max(x.chunks[axis]) * chunksize_tolerance:
        # We are staying below our threshold, so don't rechunk
        return x

    max_group = reduce(mul, map(max, x.chunks))
    new_max_group = max_group / max(x.chunks[axis]) * longest_group

    new_chunks = []
    for i in range(len(x.chunks)):
        if i == axis:
            new_chunks.append(x.chunks[i])
            continue

        new_chunk = []
        # calculate what the max chunk size in this dimension is and split every
        # chunk that is larger than that
        dimension_max_chunk = max(x.chunks[i]) * (
            (max_group / new_max_group) ** (1 / (len(x.chunks) - 1))
        )
        for c in x.chunks[i]:
            if c > chunksize_tolerance * dimension_max_chunk:
                factor = c / dimension_max_chunk

Review comment (suggestion):
```suggestion
def _resize_other_dimensions(x: Array, longest_group: int, axis: int, chunks: tuple[tuple[int, ...], ...]):
```

Code after:

    The array was reordered along axis 1 according to the positional indexer that was given.

    >>> y.compute()
    array([[ 7,  6,  3,  5,  2,  4,  1,  8],
           [15, 14, 11, 13, 10, 12,  9, 16]])
    """
    if np.isnan(x.shape).any():
        raise ValueError(
            f"Shuffling only allowed with known chunk sizes. {unknown_chunk_message}"
        )
    assert isinstance(axis, int), "axis must be an integer"
    _validate_indexer(x.chunks, indexer, axis)

    x = _rechunk_other_dimensions(x, max(map(len, indexer)), axis, chunks)

    token = tokenize(x, indexer, axis)
    out_name = f"shuffle-{token}"

    chunks, layer = _shuffle(x.chunks, indexer, axis, x.name, out_name, token)
    if len(layer) == 0:
        return Array(x.dask, x.name, x.chunks, meta=x)

    graph = HighLevelGraph.from_collections(out_name, layer, dependencies=[x])

    return Array(graph, out_name, chunks, meta=x)


def _rechunk_other_dimensions(
    x: Array, longest_group: int, axis: int, chunks: Literal["auto"]
) -> Array:
    assert chunks == "auto", "Only auto is supported for now"
    chunksize_tolerance = config.get("array.shuffle.chunksize-tolerance")

    if longest_group <= max(x.chunks[axis]) * chunksize_tolerance:
        # We are staying below our threshold, so don't rechunk
        return x

    # How much the chunksizes on our shuffle axis will increase
    chunksize_inc_factor = longest_group / max(x.chunks[axis])

    new_chunks = []
    for i in range(len(x.chunks)):
        if i == axis:
            new_chunks.append(x.chunks[i])
            continue

        new_chunksizes = []
        # calculate what the max chunk size in this dimension is and split every
        # chunk that is larger than that. We split the increase factor evenly
        # between all dimensions that are not shuffled.