(*args)
| 133 | |
| 134 | |
| 135 | def unify_chunks_expr(*args): |
| 136 | # TODO(expr): This should probably be a dedicated expression |
| 137 | # This is the implementation that expects the inputs to be expressions, the public facing |
| 138 | # variant needs to sanitize the inputs |
| 139 | if not args: |
| 140 | return {}, [], False |
| 141 | arginds = list(toolz.partition(2, args)) |
| 142 | arrays, inds = zip(*arginds) |
| 143 | if all(ind is None for ind in inds): |
| 144 | return {}, list(arrays), False |
| 145 | if all(ind == inds[0] for ind in inds) and all( |
| 146 | a.chunks == arrays[0].chunks for a in arrays |
| 147 | ): |
| 148 | return dict(zip(inds[0], arrays[0].chunks)), arrays, False |
| 149 | |
| 150 | nameinds = [] |
| 151 | blockdim_dict = dict() |
| 152 | for a, ind in arginds: |
| 153 | if ind is not None and not isinstance(a, ArrayBlockwiseDep): |
| 154 | nameinds.append((a.name, ind)) |
| 155 | blockdim_dict[a.name] = a.chunks |
| 156 | else: |
| 157 | nameinds.append((a, ind)) |
| 158 | |
| 159 | chunkss = broadcast_dimensions(nameinds, blockdim_dict, consolidate=common_blockdim) |
| 160 | |
| 161 | arrays = [] |
| 162 | changed = False |
| 163 | for a, i in arginds: |
| 164 | if i is None or isinstance(a, ArrayBlockwiseDep): |
| 165 | pass |
| 166 | else: |
| 167 | chunks = tuple( |
| 168 | ( |
| 169 | chunkss[j] |
| 170 | if a.shape[n] > 1 |
| 171 | else (a.shape[n],) if not np.isnan(sum(chunkss[j])) else None |
| 172 | ) |
| 173 | for n, j in enumerate(i) |
| 174 | ) |
| 175 | if chunks != a.chunks and all(a.chunks): |
| 176 | a = a.rechunk(chunks) |
| 177 | changed = True |
| 178 | else: |
| 179 | pass |
| 180 | arrays.append(a) |
| 181 | return chunkss, arrays, changed |
| 182 | |
| 183 | |
| 184 | class Stack(ArrayExpr): |
no test coverage detected
searching dependent graphs…