A reproduction of numpy's _parse_einsum_input() which in itself is a reproduction of c side einsum parsing in python. Returns ------- input_strings : str Parsed input strings output_string : str Parsed output string operands : list of array_like
(operands)
| 29 | # See https://github.com/numpy/numpy/blob/master/LICENSE.txt |
| 30 | # or NUMPY_LICENSE.txt within this directory |
| 31 | def parse_einsum_input(operands): |
| 32 | """ |
| 33 | A reproduction of numpy's _parse_einsum_input() |
| 34 | which in itself is a reproduction of |
| 35 | c side einsum parsing in python. |
| 36 | |
| 37 | Returns |
| 38 | ------- |
| 39 | input_strings : str |
| 40 | Parsed input strings |
| 41 | output_string : str |
| 42 | Parsed output string |
| 43 | operands : list of array_like |
| 44 | The operands to use in the numpy contraction |
| 45 | Examples |
| 46 | -------- |
| 47 | The operand list is simplified to reduce printing: |
| 48 | >> a = np.random.rand(4, 4) |
| 49 | >> b = np.random.rand(4, 4, 4) |
| 50 | >> __parse_einsum_input(('...a,...a->...', a, b)) |
| 51 | ('za,xza', 'xz', [a, b]) |
| 52 | >> __parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0])) |
| 53 | ('za,xza', 'xz', [a, b]) |
| 54 | """ |
| 55 | |
| 56 | if len(operands) == 0: |
| 57 | raise ValueError("No input operands") |
| 58 | |
| 59 | if isinstance(operands[0], str): |
| 60 | subscripts = operands[0].replace(" ", "") |
| 61 | operands = [asarray(o) for o in operands[1:]] |
| 62 | |
| 63 | # Ensure all characters are valid |
| 64 | for s in subscripts: |
| 65 | if s in ".,->": |
| 66 | continue |
| 67 | if s not in einsum_symbols_set: |
| 68 | raise ValueError("Character %s is not a valid symbol." % s) |
| 69 | |
| 70 | else: |
| 71 | tmp_operands = list(operands) |
| 72 | operand_list = [] |
| 73 | subscript_list = [] |
| 74 | for _ in range(len(operands) // 2): |
| 75 | operand_list.append(tmp_operands.pop(0)) |
| 76 | subscript_list.append(tmp_operands.pop(0)) |
| 77 | |
| 78 | output_list = tmp_operands[-1] if len(tmp_operands) else None |
| 79 | operands = [asarray(v) for v in operand_list] |
| 80 | subscripts = "" |
| 81 | last = len(subscript_list) - 1 |
| 82 | for num, sub in enumerate(subscript_list): |
| 83 | for s in sub: |
| 84 | if s is Ellipsis: |
| 85 | subscripts += "..." |
| 86 | elif isinstance(s, int): |
| 87 | subscripts += einsum_symbols[s] |
| 88 | else: |