Repository: dgasmith/opt_einsum
Pull request #200 (2022-09-26 21:19:26)

[BUG FIX] Add shapes support to parse_einsum_input
## Description
This PR adds support for the shapes kwarg for `parse_einsum_input`. This also fixes a bug where `contract_path` takes in the shapes parameter but `parse_einsum_input` is called without respecting the shapes parameter.

The bug results because even though the user had input shapes as True, `contract_path` will call `parse_einsum_input` _before_ caring about shapes, and within `parse_einsum_input`, will call .shape on an array that is already a shape.

## Todos
Notable points that this PR has either accomplished or will accomplish.
  - [x] Add support for the shapes kwarg in parse_einsum_input
  - [x] Pass along the shapes kwarg in the call from contract_path

## Questions
- [x] The lint failure looks to be infra flakiness--how can I rerun?

## Status
- [x] Ready to go

codecov[bot]: # [Codecov](https://codecov.io/gh/dgasmith/opt_einsum/pull/200?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Daniel+Smith) Report
> Merging [#200](https://codecov.io/gh/dgasmith/opt_einsum/pull/200?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Daniel+Smith) (4caffe8) into [master](https://codecov.io/gh/dgasmith/opt_einsum/commit/3824e4f41db5bd0d598a1ad2a5872c74660b2666?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Daniel+Smith) (3824e4f) will **decrease** coverage by `3.25%`.
> The diff coverage is `92.30%`.

> :exclamation: Current head 4caffe8 differs from pull request most recent head 265aefa. Consider uploading reports for the commit 265aefa to get more accurate results

dgasmith: > Welp, lol, messing with the parameters did not go well--I'll take another look tomorrow but let me know if you have any advice--I kept getting errors when I didn't use *operands about shapes getting multiple inputs :c.

That is very odd- can you post the exact error message?
janeyx99: > That is very odd- can you post the exact error message?

We are good now!
janeyx99: Is there anything on my end I should do to merge this PR?
dgasmith: I'll recap from a comment above:
> We have a faulty test in this case, It just so happens that len(arr) provides the shape of the first dimension. We need to ensure that a corresponding test fails for the given test case. A fairly minor point, but the tests as written can be misleading in the future.

I think we need to make sure the test properly raises if an array is passed as a shape. Either that or detect automatically if the object has a `.shape` attribute or not.
janeyx99: Ah oops, I didn't see that last comment. I've added a test case now.
janeyx99: Lint failure looks unrelated
dgasmith: Looks like MyPy bumped, would you mind fixing? Removing the `"type: ignore"` from line `opt_einsum/paths.py:1251` should patch it.

Patch:
diff --git a/README.md b/README.md
index 98f5d16..2258cfa 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ information.
 
 ## Example usage
 
-The [`opt_einsum.contract`](https://dgasmith.github.io/opt_einsum/api_reference.html#opt_einsumcontract)
+The [`opt_einsum.contract`](https://dgasmith.github.io/opt_einsum/api_reference#opt_einsumcontract)
 function can often act as a drop-in replacement for `einsum`
 functions without further changes to the code while providing superior performance.
 Here, a tensor contraction is performed with and without optimization:
@@ -50,7 +50,7 @@ I = np.random.rand(N, N, N, N)
 
 In this particular example, we see a ~3000x performance improvement which is
 not uncommon when compared against unoptimized contractions. See the [backend
-examples](https://dgasmith.github.io/opt_einsum/getting_started/backends.html)
+examples](https://dgasmith.github.io/opt_einsum/getting_started/backends)
 for more information on using other backends.
 
 ## Features
@@ -63,19 +63,19 @@ this repository often has more up to date algorithms for complex contractions.
 
 The following capabilities are enabled by `opt_einsum`:
 
-* Inspect [detailed information](https://dgasmith.github.io/opt_einsum/paths/introduction.html) about the path chosen.
-* Perform contractions with [numerous backends](https://dgasmith.github.io/opt_einsum/getting_started/backends.html), including on the GPU and with libraries such as [TensorFlow](https://www.tensorflow.org) and [PyTorch](https://pytorch.org).
-* Generate [reusable expressions](https://dgasmith.github.io/opt_einsum/getting_started/reusing_paths.html), potentially with [constant tensors](https://dgasmith.github.io/opt_einsum/getting_started/reusing_paths.html#specifying-constants), that can be compiled for greater performance.
-* Use an arbitrary number of indices to find contractions for [hundreds or even thousands of tensors](https://dgasmith.github.io/opt_einsum/examples/large_expr_with_greedy.html).
-* Share [intermediate computations](https://dgasmith.github.io/opt_einsum/getting_started/sharing_intermediates.html) among multiple contractions.
+* Inspect [detailed information](https://dgasmith.github.io/opt_einsum/paths/introduction) about the path chosen.
+* Perform contractions with [numerous backends](https://dgasmith.github.io/opt_einsum/getting_started/backends), including on the GPU and with libraries such as [TensorFlow](https://www.tensorflow.org) and [PyTorch](https://pytorch.org).
+* Generate [reusable expressions](https://dgasmith.github.io/opt_einsum/getting_started/reusing_paths), potentially with [constant tensors](https://dgasmith.github.io/opt_einsum/getting_started/reusing_paths#specifying-constants), that can be compiled for greater performance.
+* Use an arbitrary number of indices to find contractions for [hundreds or even thousands of tensors](https://dgasmith.github.io/opt_einsum/examples/large_expr_with_greedy).
+* Share [intermediate computations](https://dgasmith.github.io/opt_einsum/getting_started/sharing_intermediates) among multiple contractions.
 * Compute gradients of tensor contractions using [autograd](https://github.com/HIPS/autograd) or [jax](https://github.com/google/jax)
 
-Please see the [documentation](https://dgasmith.github.io/opt_einsum/index.html) for more features!
+Please see the [documentation](https://dgasmith.github.io/opt_einsum/index) for more features!
 
 ## Installation
 
 `opt_einsum` can either be installed via `pip install opt_einsum` or from conda `conda install opt_einsum -c conda-forge`.
-See the installation [documentation](https://dgasmith.github.io/opt_einsum/getting_started/install.html) for further methods.
+See the installation [documentation](https://dgasmith.github.io/opt_einsum/getting_started/install) for further methods.
 
 ## Citation
 
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 0f70379..2769c2c 100644
--- a/opt_einsum/contract.py
+++ b/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(",")
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index 8e3a00b..1f3daf8 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -3,7 +3,7 @@ A functionally equivalent parser of the numpy.einsum input parser
 """
 
 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
diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py
index ac8ccd5..0e4b587 100644
--- a/opt_einsum/paths.py
+++ b/opt_einsum/paths.py
@@ -1248,7 +1248,7 @@ class DynamicProgramming(PathOptimizer):
             # the set of n tensors is represented by a bitmap: if bit j is 1,
             # tensor j is in the set, e.g. 0b100101 = {0,2,5}; set unions
             # (intersections) can then be computed by bitwise or (and);
-            x: List[Any] = [None] * 2 + [dict() for j in range(len(g) - 1)]  # type: ignore
+            x: List[Any] = [None] * 2 + [dict() for j in range(len(g) - 1)]
             x[1] = OrderedDict((1 << j, (inputs[j], 0, inputs_contractions[j])) for j in g)
 
             # convert set of tensors g to a bitmap set:


Test patch:
diff --git a/opt_einsum/tests/test_parser.py b/opt_einsum/tests/test_parser.py
index e81fafd..28877d1 100644
--- a/opt_einsum/tests/test_parser.py
+++ b/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)