Fallback implementation that supports summing an index over > 2 inputs.
(equation, *inputs)
| 567 | |
| 568 | |
| 569 | def _exponential_space_einsum(equation, *inputs): |
| 570 | """Fallback implementation that supports summing an index over > 2 inputs.""" |
| 571 | inputs = list(inputs) |
| 572 | input_shapes = [x.get_shape() for x in inputs] |
| 573 | idx_in, idx_out = _einsum_parse_and_resolve_equation(equation, input_shapes) |
| 574 | |
| 575 | idx_all = set(''.join(idx_in) + idx_out) |
| 576 | indices = ''.join(sorted(idx_all)) |
| 577 | |
| 578 | missing_idx = set(idx_out).difference(idx_all) |
| 579 | if missing_idx: |
| 580 | raise ValueError('Unknown output axes: %s' % missing_idx) |
| 581 | |
| 582 | axis_order = {} |
| 583 | for ax in indices: |
| 584 | if ax not in idx_out: |
| 585 | axis_order[ax] = len(axis_order) |
| 586 | for ax in idx_out: |
| 587 | axis_order[ax] = len(axis_order) |
| 588 | |
| 589 | # transpose inputs so axes are in order |
| 590 | for i, (input_, axes_) in enumerate(zip(inputs, idx_in)): |
| 591 | if input_.get_shape().ndims != len(axes_): |
| 592 | raise ValueError( |
| 593 | 'Input %d with axes %s has incorrect' \ |
| 594 | ' number of dimensions (expected %d, got %d)' % ( |
| 595 | i, axes_, len(axes_), input_.get_shape().ndims |
| 596 | ) |
| 597 | ) |
| 598 | |
| 599 | sorted_idx = sorted(axes_, key=axis_order.get) |
| 600 | |
| 601 | if len(set(axes_)) != len(axes_): |
| 602 | raise ValueError( |
| 603 | 'Subscript not supported: an axis appears more than once: %s' % axes_) |
| 604 | |
| 605 | if list(axes_) != sorted_idx: |
| 606 | permuted = [axes_.find(ax) for ax in sorted_idx] |
| 607 | inputs[i] = array_ops.transpose(input_, permuted) |
| 608 | idx_in[i] = sorted_idx |
| 609 | |
| 610 | reduction_idx = [] |
| 611 | shapes = [[dim if dim else -1 |
| 612 | for dim in tensor.get_shape().as_list()] |
| 613 | for tensor in inputs] |
| 614 | |
| 615 | # validate shapes for broadcasting |
| 616 | for j, ax in enumerate(sorted(idx_all, key=axis_order.get)): |
| 617 | dims = [] |
| 618 | for i, idx in enumerate(idx_in): |
| 619 | if ax not in idx: |
| 620 | shapes[i].insert(j, 1) |
| 621 | else: |
| 622 | dim = shapes[i][j] |
| 623 | if isinstance(dim, int) and dim > 1: |
| 624 | dims.append(dim) |
| 625 | |
| 626 | if len(set(dims)) > 1: |
no test coverage detected