Merge pull request #41 from fritzo/torch-rename

Rename symbols to support torch.einsum

opt_einsum/backends/torch.py

@@ -5,8 +5,7 @@
 from __future__ import absolute_import
 import numpy as np
 
-from ..parser import einsum_symbols_base
-
+from ..parser import convert_to_valid_einsum_chars, einsum_symbols_base
 
 _TORCH_DEVICE = None
 
@@ -31,6 +30,10 @@ def transpose(a, axes):
 def einsum(equation, *operands):
     """Variadic version of torch.einsum to match numpy api.
     """
+    # rename symbols to support PyTorch 0.4.1 and earlier,
+    # which allow only symbols a-z.
+    equation = convert_to_valid_einsum_chars(equation)
+
     torch, _ = _get_torch_and_device()
     return torch.einsum(equation, operands)
 
@@ -39,8 +42,6 @@ def tensordot(x, y, axes=2):
     """Simple translation of tensordot syntax to einsum.
     """
     # XXX: tensordot should be directly implemented in torch soon
-    torch, _ = _get_torch_and_device()
-
     xnd = x.ndimension()
     ynd = y.ndimension()
 
@@ -80,7 +81,7 @@ def tensordot(x, y, axes=2):
 
     # form full string and contract!
     einsum_str = "{},{}->{}".format(*map("".join, (x_ix, y_ix, out_ix)))
-    return torch.einsum(einsum_str, (x, y))
+    return einsum(einsum_str, x, y)
 
 
 def to_torch(array):

opt_einsum/helpers.py

@@ -4,8 +4,8 @@
 
 import numpy as np
 
-chars = 'abcdefghijklmopq'
-sizes = np.array([2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3, 2, 5, 7, 4, 3])
+chars = 'abcdefghijklmopqABC'
+sizes = np.array([2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3, 2, 5, 7, 4, 3, 2, 3, 4])
 default_dim_dict = {c: s for c, s in zip(chars, sizes)}
 
 

opt_einsum/parser.py

@@ -57,18 +57,11 @@ def gen_unused_symbols(used, n):
 
 def convert_to_valid_einsum_chars(einsum_str):
     """Convert the str ``einsum_str`` to contain only the alphabetic characters
-    valid for numpy einsum.
+    valid for numpy einsum. If there are too many symbols, let the backend
+    throw an error.
     """
-    # partition into valid and invalid sets
-    valid, invalid = set(), set()
-    for x in einsum_str:
-        (valid if is_valid_einsum_char(x) else invalid).add(x)
-
-    # get replacements for invalid chars that are not already used
-    available = gen_unused_symbols(valid, len(invalid))
-
-    # map invalid to available and replace in the inputs
-    replacer = dict(zip(invalid, available))
+    symbols = sorted(set(einsum_str) - set(',->'))
+    replacer = {x: get_symbol(i) for i, x in enumerate(symbols)}
     return "".join(replacer.get(x, x) for x in einsum_str)
 
 

opt_einsum/tests/test_backends.py

@@ -35,6 +35,7 @@
     'ijk,ikj',
     'i,j->ij',
     'ijk,k->ij',
+    'AB,BC->CA',
 ]