| 361 | |
| 362 | |
| 363 | def _chunk_sum(a, axis=None, dtype=None, keepdims=None): |
| 364 | # Caution: this is not your conventional array-sum: due |
| 365 | # to the special nature of the preceding blockwise con- |
| 366 | # traction, each chunk is expected to have exactly the |
| 367 | # same shape, with a size of 1 for the dimension given |
| 368 | # by `axis` (the reduction axis). This makes mere ele- |
| 369 | # ment-wise addition of the arrays possible. Besides, |
| 370 | # the output can be merely squeezed to lose the `axis`- |
| 371 | # dimension when keepdims = False |
| 372 | if type(a) is list: |
| 373 | out = reduce(partial(np.add, dtype=dtype), a) |
| 374 | else: |
| 375 | out = a |
| 376 | |
| 377 | if keepdims: |
| 378 | return out |
| 379 | else: |
| 380 | return out.squeeze(axis[0]) |
| 381 | |
| 382 | |
| 383 | def _sum_wo_cat(a, axis=None, dtype=None): |