[BUG FIX] Add shapes support to parse_einsum_input (#200)

* [BUG FIX] Add shapes support to parse_einsum_input

* Realized more than just having shape is important for converting to numpy

* fix lint

* Address comments, add test case

* Just smoosh args into a list

* Make the test make sense

* Raise error when non-shapes look passed in but shapes=True

* Patch mypy, use pytest

Author: janeyx99

opt_einsum/contract.py

@@ -129,7 +129,7 @@ def contract_path(*operands_: Any, **kwargs: Any) -> Tuple[PathType, PathInfo]:
     **Parameters:**
 
     - **subscripts** - *(str)* Specifies the subscripts for summation.
-    - **\\*operands** - *(list of array_like)* hese are the arrays for the operation.
+    - **\\*operands** - *(list of array_like)* these are the arrays for the operation.
     - **use_blas** - *(bool)* Do you use BLAS for valid operations, may use extra memory for more intermediates.
     - **optimize** - *(str, list or bool, optional (default: `auto`))* Choose the type of path.
 
@@ -251,7 +251,7 @@ def contract_path(*operands_: Any, **kwargs: Any) -> Tuple[PathType, PathInfo]:
     use_blas = kwargs.pop("use_blas", True)
 
     # Python side parsing
-    input_subscripts, output_subscript, operands = parser.parse_einsum_input(operands_)
+    input_subscripts, output_subscript, operands = parser.parse_einsum_input(operands_, shapes=shapes)
 
     # Build a few useful list and sets
     input_list = input_subscripts.split(",")

opt_einsum/parser.py

@@ -3,7 +3,7 @@
 """
 
 import itertools
-from typing import Any, Dict, Iterator, List, Tuple
+from typing import Any, Dict, Iterator, List, Tuple, Union
 
 import numpy as np
 
@@ -224,7 +224,7 @@ def convert_subscripts(old_sub: List[Any], symbol_map: Dict[Any, Any]) -> str:
     return new_sub
 
 
-def convert_interleaved_input(operands: List[Any]) -> Tuple[str, List[Any]]:
+def convert_interleaved_input(operands: Union[List[Any], Tuple[Any]]) -> Tuple[str, List[Any]]:
     """Convert 'interleaved' input to standard einsum input."""
     tmp_operands = list(operands)
     operand_list = []
@@ -262,10 +262,16 @@ def convert_interleaved_input(operands: List[Any]) -> Tuple[str, List[Any]]:
     return subscripts, operands
 
 
-def parse_einsum_input(operands: Any) -> Tuple[str, str, List[ArrayType]]:
+def parse_einsum_input(operands: Any, shapes: bool = False) -> Tuple[str, str, List[ArrayType]]:
     """
     A reproduction of einsum c side einsum parsing in python.
 
+    **Parameters:**
+    Intakes the same inputs as `contract_path`, but NOT the keyword args. The only
+    supported keyword argument is:
+    - **shapes** - *(bool, optional)* Whether ``parse_einsum_input`` should assume arrays (the default) or
+        array shapes have been supplied.
+
     Returns
     -------
     input_strings : str
@@ -293,11 +299,21 @@ def parse_einsum_input(operands: Any) -> Tuple[str, str, List[ArrayType]]:
 
     if isinstance(operands[0], str):
         subscripts = operands[0].replace(" ", "")
+        if shapes:
+            if any([hasattr(o, "shape") for o in operands[1:]]):
+                raise ValueError(
+                    "shapes is set to True but given at least one operand looks like an array"
+                    " (at least one operand has a shape attribute). "
+                )
         operands = [possibly_convert_to_numpy(x) for x in operands[1:]]
-
     else:
         subscripts, operands = convert_interleaved_input(operands)
 
+    if shapes:
+        operand_shapes = operands
+    else:
+        operand_shapes = [o.shape for o in operands]
+
     # Check for proper "->"
     if ("-" in subscripts) or (">" in subscripts):
         invalid = (subscripts.count("-") > 1) or (subscripts.count(">") > 1)
@@ -307,7 +323,7 @@ def parse_einsum_input(operands: Any) -> Tuple[str, str, List[ArrayType]]:
     # Parse ellipses
     if "." in subscripts:
         used = subscripts.replace(".", "").replace(",", "").replace("->", "")
-        ellipse_inds = "".join(gen_unused_symbols(used, max(len(x.shape) for x in operands)))
+        ellipse_inds = "".join(gen_unused_symbols(used, max(len(x) for x in operand_shapes)))
         longest = 0
 
         # Do we have an output to account for?
@@ -325,10 +341,10 @@ def parse_einsum_input(operands: Any) -> Tuple[str, str, List[ArrayType]]:
                     raise ValueError("Invalid Ellipses.")
 
                 # Take into account numerical values
-                if operands[num].shape == ():
+                if operand_shapes[num] == ():
                     ellipse_count = 0
                 else:
-                    ellipse_count = max(len(operands[num].shape), 1) - (len(sub) - 3)
+                    ellipse_count = max(len(operand_shapes[num]), 1) - (len(sub) - 3)
 
                 if ellipse_count > longest:
                     longest = ellipse_count
@@ -370,6 +386,9 @@ def parse_einsum_input(operands: Any) -> Tuple[str, str, List[ArrayType]]:
 
     # Make sure number operands is equivalent to the number of terms
     if len(input_subscripts.split(",")) != len(operands):
-        raise ValueError("Number of einsum subscripts must be equal to the " "number of operands.")
+        raise ValueError(
+            f"Number of einsum subscripts, {len(input_subscripts.split(','))}, must be equal to the "
+            f"number of operands, {len(operands)}."
+        )
 
     return input_subscripts, output_subscript, operands

opt_einsum/tests/test_parser.py

@@ -2,7 +2,10 @@
 Directly tests various parser utility functions.
 """
 
-from opt_einsum.parser import get_symbol
+from multiprocessing.sharedctypes import Value
+import numpy as np
+import pytest
+from opt_einsum.parser import get_symbol, parse_einsum_input, possibly_convert_to_numpy
 
 
 def test_get_symbol():
@@ -12,3 +15,29 @@ def test_get_symbol():
     assert get_symbol(55295) == "\ud88b"
     assert get_symbol(55296) == "\ue000"
     assert get_symbol(57343) == "\ue7ff"
+
+
+def test_parse_einsum_input():
+    eq = "ab,bc,cd"
+    ops = [np.random.rand(2, 3), np.random.rand(3, 4), np.random.rand(4, 5)]
+    input_subscripts, output_subscript, operands = parse_einsum_input([eq, *ops])
+    assert input_subscripts == eq
+    assert output_subscript == "ad"
+    assert operands == ops
+
+
+def test_parse_einsum_input_shapes_error():
+    eq = "ab,bc,cd"
+    ops = [np.random.rand(2, 3), np.random.rand(3, 4), np.random.rand(4, 5)]
+
+    with pytest.raises(ValueError):
+        _ = parse_einsum_input([eq, *ops], shapes=True)
+
+
+def test_parse_einsum_input_shapes():
+    eq = "ab,bc,cd"
+    shps = [(2, 3), (3, 4), (4, 5)]
+    input_subscripts, output_subscript, operands = parse_einsum_input([eq, *shps], shapes=True)
+    assert input_subscripts == eq
+    assert output_subscript == "ad"
+    assert np.allclose([possibly_convert_to_numpy(shp) for shp in shps], operands)