Repository: dgasmith/opt_einsum (python)

Issue:
Support for get_symbol(116) in Python 2
`get_symbol(x)` currently fails in Python 2 for any `x >= 116`:
```
i = 116

    def get_symbol(i):
        """Get the symbol corresponding to int ``i`` - runs through the usual 52
        letters before resorting to unicode characters, starting at ``chr(192)``.

        Examples
        --------
        >>> get_symbol(2)
        'c'

        >>> oe.get_symbol(200)
        'Ŕ'

        >>> oe.get_symbol(20000)
        '京'
        """
        if i < 52:
            return einsum_symbols_base[i]
>       return chr(i + 140)
E       ValueError: chr() arg not in range(256)
```
Can we fix this by use `unichr()` and unicode strings in Python 2?

Pull request description:
Make get_symbol(large_number) work in Python2
Fixes #42

## Description
This provides a compatibility shim for Python 2 so that `get_symbol(n)` can work for larger n.

Note that Python 2's `numpy.einsum` still does not accept unicode. This means that, while `opt_einsum` can now work with arbitrarily large contractions in Python 2, it only works with contractions whose individual calls to Numpy involve fewer than 52 (I believe) characters.

## Tested

- [x] added a regression test
- [x] enabled previously skipped tests
- [x] tested in a Pyro application

## Status
- [x] Ready to go

Patch:
diff --git a/opt_einsum/compat.py b/opt_einsum/compat.py
new file mode 100644
index 0000000..38cf55d
--- /dev/null
+++ b/opt_einsum/compat.py
@@ -0,0 +1,10 @@
+# Python 2/3 compatability shim
+
+try:
+    # Python 2
+    get_chr = unichr
+    strings = (str, type(get_chr(300)))
+except NameError:
+    # Python 3
+    get_chr = chr
+    strings = str
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index a7932f0..01998dd 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -6,6 +6,7 @@ import numpy as np
 
 from . import backends
 from . import blas
+from . import compat
 from . import helpers
 from . import parser
 from . import paths
@@ -176,7 +177,7 @@ def contract_path(*operands, **kwargs):
     naive_cost = helpers.flop_count(indices, inner_product, num_ops, dimension_dict)
 
     # Compute the path
-    if not isinstance(path_type, str):
+    if not isinstance(path_type, compat.strings):
         path = path_type
     elif num_ops == 1:
         # Nothing to be optimized
@@ -274,7 +275,7 @@ def _einsum(*operands, **kwargs):
     """
     fn = backends.get_func('einsum', kwargs.pop('backend', 'numpy'))
 
-    if not isinstance(operands[0], str):
+    if not isinstance(operands[0], compat.strings):
         return fn(*operands, **kwargs)
 
     einsum_str, operands = operands[0], operands[1:]
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index c5370b1..adfc19f 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -8,6 +8,8 @@ from collections import OrderedDict
 
 import numpy as np
 
+from . import compat
+
 einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 
 
@@ -40,7 +42,7 @@ def get_symbol(i):
     """
     if i < 52:
         return einsum_symbols_base[i]
-    return chr(i + 140)
+    return compat.get_chr(i + 140)
 
 
 def gen_unused_symbols(used, n):
@@ -133,7 +135,7 @@ def parse_einsum_input(operands):
     if len(operands) == 0:
         raise ValueError("No input operands")
 
-    if isinstance(operands[0], str):
+    if isinstance(operands[0], compat.strings):
         subscripts = operands[0].replace(" ", "")
         operands = [possibly_convert_to_numpy(x) for x in operands[1:]]
 
diff --git a/opt_einsum/sharing.py b/opt_einsum/sharing.py
index 504ba05..54c94bd 100644
--- a/opt_einsum/sharing.py
+++ b/opt_einsum/sharing.py
@@ -129,7 +129,7 @@ def einsum_cache_wrap(einsum):
         canonical = sorted(zip(inputs, map(id, operands)), key=lambda x: x[1])
         canonical_ids = tuple(id_ for _, id_ in canonical)
         canonical_inputs = ','.join(input_ for input_, _ in canonical)
-        canonical_equation = alpha_canonicalize('{}->{}'.format(canonical_inputs, output))
+        canonical_equation = alpha_canonicalize(canonical_inputs + "->" + output)
         key = 'einsum', backend, canonical_equation, canonical_ids
         return _memoize(key, einsum, equation, *operands, backend=backend)
 


Test patch:
diff --git a/opt_einsum/tests/test_contract.py b/opt_einsum/tests/test_contract.py
index 43b5061..fdfacc7 100644
--- a/opt_einsum/tests/test_contract.py
+++ b/opt_einsum/tests/test_contract.py
@@ -7,7 +7,7 @@ import sys
 import numpy as np
 import pytest
 
-from opt_einsum import contract, contract_path, helpers, contract_expression
+from opt_einsum import compat, contract, contract_path, helpers, contract_expression
 
 tests = [
     # Test hadamard-like products
@@ -117,7 +117,6 @@ def test_drop_in_replacement(string):
     assert np.allclose(opt, np.einsum(string, *views))
 
 
-@pytest.mark.skipif(sys.version_info[0] < 3, reason='requires python3')
 @pytest.mark.parametrize("string", tests)
 def test_compare_greek(string):
     views = helpers.build_views(string)
@@ -125,7 +124,7 @@ def test_compare_greek(string):
     ein = contract(string, *views, optimize=False, use_blas=False)
 
     # convert to greek
-    string = ''.join(chr(ord(c) + 848) if c not in ',->.' else c for c in string)
+    string = ''.join(compat.get_chr(ord(c) + 848) if c not in ',->.' else c for c in string)
 
     opt = contract(string, *views, optimize='greedy', use_blas=False)
     assert np.allclose(ein, opt)
@@ -146,7 +145,6 @@ def test_compare_blas(string):
     assert np.allclose(ein, opt)
 
 
-@pytest.mark.skipif(sys.version_info[0] < 3, reason='requires python3')
 @pytest.mark.parametrize("string", tests)
 def test_compare_blas_greek(string):
     views = helpers.build_views(string)
@@ -154,7 +152,7 @@ def test_compare_blas_greek(string):
     ein = contract(string, *views, optimize=False)
 
     # convert to greek
-    string = ''.join(chr(ord(c) + 848) if c not in ',->.' else c for c in string)
+    string = ''.join(compat.get_chr(ord(c) + 848) if c not in ',->.' else c for c in string)
 
     opt = contract(string, *views, optimize='greedy')
     assert np.allclose(ein, opt)
@@ -163,10 +161,9 @@ def test_compare_blas_greek(string):
     assert np.allclose(ein, opt)
 
 
-@pytest.mark.skipif(sys.version_info[0] < 3, reason='requires python3')
 def test_some_non_alphabet_maintains_order():
     # 'c beta a' should automatically go to -> 'a c beta'
-    string = 'c' + chr(ord('b') + 848) + 'a'
+    string = 'c' + compat.get_chr(ord('b') + 848) + 'a'
     # but beta will be temporarily replaced with 'b' for which 'cba->abc'
     # so check manual output kicks in:
     x = np.random.rand(2, 3, 4)
diff --git a/opt_einsum/tests/test_input.py b/opt_einsum/tests/test_input.py
index 022a078..311356e 100644
--- a/opt_einsum/tests/test_input.py
+++ b/opt_einsum/tests/test_input.py
@@ -230,7 +230,6 @@ def test_singleton_dimension_broadcast():
         assert np.allclose(res2, np.full((1, 5), 5))
 
 
-@pytest.mark.skipif(sys.version_info.major == 2, reason="Requires python 3.")
 def test_large_int_input_format():
     string = 'ab,bc,cd'
     x, y, z = build_views(string)
diff --git a/opt_einsum/tests/test_paths.py b/opt_einsum/tests/test_paths.py
index 8087dd2..da39e23 100644
--- a/opt_einsum/tests/test_paths.py
+++ b/opt_einsum/tests/test_paths.py
@@ -3,6 +3,8 @@ Tests the accuracy of the opt_einsum paths in addition to unit tests for
 the various path helper functions.
 """
 
+import itertools
+
 import numpy as np
 import pytest
 
@@ -174,3 +176,14 @@ def test_can_optimize_outer_products():
     a, b, c = [np.random.randn(10, 10) for _ in range(3)]
     d = np.random.randn(10, 2)
     assert oe.contract_path("ab,cd,ef,fg", a, b, c, d, path='greedy')[0] == [(2, 3), (0, 2), (0, 1)]
+
+
+@pytest.mark.parametrize('num_symbols', [2, 3, 26, 26 + 26, 256 - 140, 300])
+def test_large_path(num_symbols):
+    symbols = ''.join(oe.get_symbol(i) for i in range(num_symbols))
+    dimension_dict = dict(zip(symbols, itertools.cycle([2, 3, 4])))
+    expression = ','.join(symbols[t:t+2] for t in range(num_symbols - 1))
+    tensors = oe.helpers.build_views(expression, dimension_dict=dimension_dict)
+
+    # Check that path construction does not crash
+    oe.contract_path(expression, *tensors, path='greedy')