(
nd_v_chunks: tuple[tuple[int, ...], ...],
nd_backend_chunks: tuple[tuple[int, ...], ...],
)
| 4 | |
| 5 | |
| 6 | def align_nd_chunks( |
| 7 | nd_v_chunks: tuple[tuple[int, ...], ...], |
| 8 | nd_backend_chunks: tuple[tuple[int, ...], ...], |
| 9 | ) -> tuple[tuple[int, ...], ...]: |
| 10 | if len(nd_backend_chunks) != len(nd_v_chunks): |
| 11 | raise ValueError( |
| 12 | "The number of dimensions on the backend and the variable must be the same." |
| 13 | ) |
| 14 | |
| 15 | nd_aligned_chunks: list[tuple[int, ...]] = [] |
| 16 | for backend_chunks, v_chunks in zip(nd_backend_chunks, nd_v_chunks, strict=True): |
| 17 | # Validate that they have the same number of elements |
| 18 | if sum(backend_chunks) != sum(v_chunks): |
| 19 | raise ValueError( |
| 20 | "The number of elements in the backend does not " |
| 21 | "match the number of elements in the variable. " |
| 22 | "This inconsistency should never occur at this stage." |
| 23 | ) |
| 24 | |
| 25 | # Validate if the backend_chunks satisfy the condition that all the values |
| 26 | # excluding the borders are equal |
| 27 | if len(set(backend_chunks[1:-1])) > 1: |
| 28 | raise ValueError( |
| 29 | f"This function currently supports aligning chunks " |
| 30 | f"only when backend chunks are of uniform size, excluding borders. " |
| 31 | f"If you encounter this error, please report it—this scenario should never occur " |
| 32 | f"unless there is an internal misuse. " |
| 33 | f"Backend chunks: {backend_chunks}" |
| 34 | ) |
| 35 | |
| 36 | # The algorithm assumes that there are always two borders on the |
| 37 | # Backend and the Array if not, the result is going to be the same |
| 38 | # as the input, and there is nothing to optimize |
| 39 | if len(backend_chunks) == 1: |
| 40 | nd_aligned_chunks.append(backend_chunks) |
| 41 | continue |
| 42 | |
| 43 | if len(v_chunks) == 1: |
| 44 | nd_aligned_chunks.append(v_chunks) |
| 45 | continue |
| 46 | |
| 47 | # Size of the chunk on the backend |
| 48 | fixed_chunk = max(backend_chunks) |
| 49 | |
| 50 | # The ideal size of the chunks is the maximum of the two; this would avoid |
| 51 | # that we use more memory than expected |
| 52 | max_chunk = max(fixed_chunk, *v_chunks) |
| 53 | |
| 54 | # The algorithm assumes that the chunks on this array are aligned except the last one |
| 55 | # because it can be considered a partial one |
| 56 | aligned_chunks: list[int] = [] |
| 57 | |
| 58 | # For simplicity of the algorithm, let's transform the Array chunks in such a way that |
| 59 | # we remove the partial chunks. To achieve this, we add artificial data to the borders |
| 60 | t_v_chunks = list(v_chunks) |
| 61 | t_v_chunks[0] += fixed_chunk - backend_chunks[0] |
| 62 | t_v_chunks[-1] += fixed_chunk - backend_chunks[-1] |
| 63 |
no outgoing calls
searching dependent graphs…