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

Code before:
    _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
                nc = [c // math.ceil(factor)] * math.ceil(factor)
                for ii in range(c - sum(nc)):
                    nc[ii] += 1
                new_chunk.extend(nc)
            else:
                new_chunk.append(c)
        new_chunks.append(new_chunk)

    return x.rechunk(tuple(map(tuple, new_chunks)))

Review comment (suggestion):
```suggestion
    max_axis_chunks = max(x.chunks[axis])
```

Code after:
            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.
        up_chunksize_limit_for_dim = max(x.chunks[i]) / (
            chunksize_inc_factor ** (1 / (len(x.chunks) - 1))
        )
        for c in x.chunks[i]:
            if c > chunksize_tolerance * up_chunksize_limit_for_dim:
                factor = math.ceil(c / up_chunksize_limit_for_dim)

                # Ensure that we end up at least with chunksize 1
                factor = min(factor, c)