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: 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.
        up_chunksize_limit_for_dim = max(x.chunks[i]) / (
            chunksize_inc_factor ** (1 / (len(x.chunks) - 1))
        )

Review comment (suggestion):
```suggestion
def _rechunk_other_dimensions(
```

Code after:
    ((2,), (5, 3))

    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