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

Code before:
    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.chunk-size-tolerance")

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

    # How large is the largest chunk in the input
    maximum_chunk = reduce(mul, map(max, x.chunks))

    changeable_dimensions = set(range(len(x.chunks))) - {axis}
    new_chunks = list(x.chunks)
    new_chunks[axis] = (longest_group,)

    # iterate until we distributed the increase in chunksize accross all dimensions
    # or every non-shuffle dimension is all 1
    while changeable_dimensions:
        length_changeable_dimensions = len(changeable_dimensions)
        chunksize_inc_factor = reduce(mul, map(max, new_chunks)) / maximum_chunk  # type: ignore[operator]
        if chunksize_inc_factor <= 1:
            break

        for i in list(changeable_dimensions):
            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(new_chunks[i]) / (
                chunksize_inc_factor ** (1 / length_changeable_dimensions)
            )
            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)

                    chunksize, remainder = divmod(c, factor)
                    nc = [chunksize] * factor
                    for ii in range(remainder):
                        # Add remainder parts to the first few chunks
                        nc[ii] += 1
                    new_chunksizes.extend(nc)

Review comment (suggestion):
super nit:
```suggestion
        n_changeable_dimensions = len(changeable_dimensions)
```

Code after:
    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.chunk-size-tolerance")

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

    # How large is the largest chunk in the input
    maximum_chunk = reduce(mul, map(max, x.chunks))

    changeable_dimensions = set(range(len(x.chunks))) - {axis}
    new_chunks = list(x.chunks)
    new_chunks[axis] = (longest_group,)

    # iterate until we distributed the increase in chunksize accross all dimensions
    # or every non-shuffle dimension is all 1
    while changeable_dimensions:
        n_changeable_dimensions = len(changeable_dimensions)
        chunksize_inc_factor = reduce(mul, map(max, new_chunks)) / maximum_chunk  # type: ignore[operator]
        if chunksize_inc_factor <= 1:
            break

        for i in list(changeable_dimensions):
            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(new_chunks[i]) / (
                chunksize_inc_factor ** (1 / n_changeable_dimensions)
            )
            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)

                    chunksize, remainder = divmod(c, factor)
                    nc = [chunksize] * factor
                    for ii in range(remainder):
                        # Add remainder parts to the first few chunks
                        nc[ii] += 1
                    new_chunksizes.extend(nc)