Given array shapes, return the resulting shape and slices prefixes. These help in nested concatenation. Returns ------- shape: tuple of int This tuple satisfies:: shape, _ = _concatenate_shapes([arr.shape for shape in arrs], axis) shape == concatena
(shapes, axis)
| 561 | |
| 562 | |
| 563 | def _concatenate_shapes(shapes, axis): |
| 564 | """Given array shapes, return the resulting shape and slices prefixes. |
| 565 | |
| 566 | These help in nested concatenation. |
| 567 | |
| 568 | Returns |
| 569 | ------- |
| 570 | shape: tuple of int |
| 571 | This tuple satisfies:: |
| 572 | |
| 573 | shape, _ = _concatenate_shapes([arr.shape for shape in arrs], axis) |
| 574 | shape == concatenate(arrs, axis).shape |
| 575 | |
| 576 | slice_prefixes: tuple of (slice(start, end), ) |
| 577 | For a list of arrays being concatenated, this returns the slice |
| 578 | in the larger array at axis that needs to be sliced into. |
| 579 | |
| 580 | For example, the following holds:: |
| 581 | |
| 582 | ret = concatenate([a, b, c], axis) |
| 583 | _, (sl_a, sl_b, sl_c) = concatenate_slices([a, b, c], axis) |
| 584 | |
| 585 | ret[(slice(None),) * axis + sl_a] == a |
| 586 | ret[(slice(None),) * axis + sl_b] == b |
| 587 | ret[(slice(None),) * axis + sl_c] == c |
| 588 | |
| 589 | These are called slice prefixes since they are used in the recursive |
| 590 | blocking algorithm to compute the left-most slices during the |
| 591 | recursion. Therefore, they must be prepended to rest of the slice |
| 592 | that was computed deeper in the recursion. |
| 593 | |
| 594 | These are returned as tuples to ensure that they can quickly be added |
| 595 | to existing slice tuple without creating a new tuple every time. |
| 596 | |
| 597 | """ |
| 598 | # Cache a result that will be reused. |
| 599 | shape_at_axis = [shape[axis] for shape in shapes] |
| 600 | |
| 601 | # Take a shape, any shape |
| 602 | first_shape = shapes[0] |
| 603 | first_shape_pre = first_shape[:axis] |
| 604 | first_shape_post = first_shape[axis+1:] |
| 605 | |
| 606 | if any(shape[:axis] != first_shape_pre or |
| 607 | shape[axis+1:] != first_shape_post for shape in shapes): |
| 608 | raise ValueError( |
| 609 | 'Mismatched array shapes in block along axis {}.'.format(axis)) |
| 610 | |
| 611 | shape = (first_shape_pre + (sum(shape_at_axis),) + first_shape[axis+1:]) |
| 612 | |
| 613 | offsets_at_axis = _accumulate(shape_at_axis) |
| 614 | slice_prefixes = [(slice(start, end),) |
| 615 | for start, end in zip([0] + offsets_at_axis, |
| 616 | offsets_at_axis)] |
| 617 | return shape, slice_prefixes |
| 618 | |
| 619 | |
| 620 | def _block_info_recursion(arrays, max_depth, result_ndim, depth=0): |
no test coverage detected