Partial reduction across multiple axes. Parameters ---------- func : function x : Array split_every : dict Maximum reduction block sizes in each dimension. Examples -------- Reduce across axis 0 and 2, merging a maximum of 1 block in the 0th dimension, a
(
func, x, split_every, keepdims=False, dtype=None, name=None, reduced_meta=None
)
| 271 | |
| 272 | |
| 273 | def partial_reduce( |
| 274 | func, x, split_every, keepdims=False, dtype=None, name=None, reduced_meta=None |
| 275 | ): |
| 276 | """Partial reduction across multiple axes. |
| 277 | |
| 278 | Parameters |
| 279 | ---------- |
| 280 | func : function |
| 281 | x : Array |
| 282 | split_every : dict |
| 283 | Maximum reduction block sizes in each dimension. |
| 284 | |
| 285 | Examples |
| 286 | -------- |
| 287 | Reduce across axis 0 and 2, merging a maximum of 1 block in the 0th |
| 288 | dimension, and 3 blocks in the 2nd dimension: |
| 289 | |
| 290 | >>> partial_reduce(np.min, x, {0: 1, 2: 3}) # doctest: +SKIP |
| 291 | """ |
| 292 | name = ( |
| 293 | (name or funcname(func)) + "-" + tokenize(func, x, split_every, keepdims, dtype) |
| 294 | ) |
| 295 | parts = [ |
| 296 | list(partition_all(split_every.get(i, 1), range(n))) |
| 297 | for (i, n) in enumerate(x.numblocks) |
| 298 | ] |
| 299 | keys = product(*map(range, map(len, parts))) |
| 300 | out_chunks = [ |
| 301 | tuple(1 for p in partition_all(split_every[i], c)) if i in split_every else c |
| 302 | for (i, c) in enumerate(x.chunks) |
| 303 | ] |
| 304 | if not keepdims: |
| 305 | out_axis = [i for i in range(x.ndim) if i not in split_every] |
| 306 | getter = lambda k: get(out_axis, k) |
| 307 | keys = map(getter, keys) |
| 308 | out_chunks = list(getter(out_chunks)) |
| 309 | dsk = {} |
| 310 | for k, p in zip(keys, product(*parts)): |
| 311 | free = { |
| 312 | i: j[0] for (i, j) in enumerate(p) if len(j) == 1 and i not in split_every |
| 313 | } |
| 314 | dummy = dict(i for i in enumerate(p) if i[0] in split_every) |
| 315 | g = lol_tuples((x.name,), range(x.ndim), free, dummy) |
| 316 | dsk[(name,) + k] = (func, g) |
| 317 | graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x]) |
| 318 | |
| 319 | meta = x._meta |
| 320 | if reduced_meta is not None: |
| 321 | try: |
| 322 | meta = func(reduced_meta, computing_meta=True) |
| 323 | # no meta keyword argument exists for func, and it isn't required |
| 324 | except TypeError: |
| 325 | try: |
| 326 | meta = func(reduced_meta) |
| 327 | except ValueError as e: |
| 328 | # min/max functions have no identity, don't apply function to meta |
| 329 | if "zero-size array to reduction operation" in str(e): |
| 330 | meta = reduced_meta |
no test coverage detected
searching dependent graphs…