check equation / raise error, default right labels generation
(
equation: str, *operands: Tensor
)
| 779 | |
| 780 | |
| 781 | def preprocess( |
| 782 | equation: str, *operands: Tensor |
| 783 | ) -> tuple[str, str, list[str], list[Tensor]]: |
| 784 | """ |
| 785 | check equation / raise error, default right labels generation |
| 786 | """ |
| 787 | equation = equation.replace(" ", "") |
| 788 | nop = len(operands) |
| 789 | assert nop > 0, ( |
| 790 | f"Required at least one operand in Einsum API, but received {nop}" |
| 791 | ) |
| 792 | |
| 793 | # Part the equation to left hand side and right hand side |
| 794 | lhs, *rhs = equation.lower().split('->') |
| 795 | assert len(rhs) < 2, "Invalid equation: multiple `->` were found." |
| 796 | |
| 797 | labels = parse_labels(lhs, operands) |
| 798 | # Note, we distinguish between 'ij->' and 'ij' by setting rhs to '' and None |
| 799 | rhs = rhs[0] if rhs else None |
| 800 | if rhs is None: |
| 801 | rhs = rhs_inference(lhs) |
| 802 | |
| 803 | assert len(lhs.split(',')) == len(operands), ( |
| 804 | f"Invalid equation: the number of operands is {len(operands)}, " |
| 805 | f"but found {len(lhs.split(','))} segments in the label equation." |
| 806 | ) |
| 807 | |
| 808 | assert not ('...' in lhs and '...' not in rhs), ( |
| 809 | 'Invalid equation: missing ellipsis in output labels.' |
| 810 | ) |
| 811 | |
| 812 | lhs, rhs, new_operands = replace_ellipsis(lhs, rhs, *operands) |
| 813 | return lhs, rhs, labels, new_operands |
| 814 | |
| 815 | |
| 816 | class Shaped(NamedTuple): |
no test coverage detected