Repository: dgasmith/opt_einsum

Issue:
reusable einsum expressions
Essentially, I'm using ``opt_einsum`` for some tensor network calculations, and a significant proportion of time is spent in ``contract_path``, even for relatively large sized tensors and pre-calculated paths.

Specifically, a huge number of contractions with the same indices and shapes are performed over and over again and I already cache these using the path from ``contract_path``. However, it seems even with the path supplied, a large amount of time is spent parsing the input, in ``can_blas``, ``parse_einsum_input`` etc.

My suggestion would be something like:

```python
shapes = [(2, 3), (3, 4), (4, 5)]
my_expr = einsum_expression("ab,bc,cd->ad", *shapes)

for _ in many_repeats:
    ...
    x, y, z = (rand(s) for s in shapes)
    out = my_expr(x, y, z)
    ...
```

I.e. it would only accept arrays of the same dimensions (and probably an ``out`` argument) and would otherwise skip all parsing, using a ``contraction_list`` stored within the function.

Anyway, I'm about to knock something up to test for myself, but thought it might be of more general interest.

Comments:

Patch:
diff --git a/README.md b/README.md
index a792df7..9e4ed20 100644
--- a/README.md
+++ b/README.md
@@ -131,6 +131,35 @@ True
 
 By contracting terms in the correct order we can see that this expression can be computed with N^4 scaling. Even with the overhead of finding the best order or 'path' and small dimensions, opt_einsum is roughly 900 times faster than pure einsum for this expression.
 
+
+## Reusing paths using ``contract_expression``
+
+If you expect to repeatedly use a particular contraction it can make things simpler and more efficient to not compute the path each time. Instead, supplying ``contract_expression`` with the contraction string and the shapes of the tensors generates a ``ContractExpression`` which can then be repeatedly called with any matching set of arrays. For example:
+
+```python
+>>> my_expr = oe.contract_expression("abc,cd,dbe->ea", (2, 3, 4), (4, 5), (5, 3, 6))
+>>> print(my_expr)
+<ContractExpression> for 'abc,cd,dbe->ea':
+  1.  'dbe,cd->bce' [GEMM]
+  2.  'bce,abc->ea' [GEMM]
+```
+
+Now we can call this expression with 3 arrays that match the original shapes without having to compute the path again:
+
+```python
+>>> x, y, z = (np.random.rand(*s) for s in [(2, 3, 4), (4, 5), (5, 3, 6)])
+>>> my_expr(x, y, z)
+array([[ 3.08331541,  4.13708916],
+       [ 2.92793729,  4.57945185],
+       [ 3.55679457,  5.56304115],
+       [ 2.6208398 ,  4.39024187],
+       [ 3.66736543,  5.41450334],
+       [ 3.67772272,  5.46727192]])
+```
+
+Note that few checks are performed when calling the expression, and while it will work for a set of arrays with the same ranks as the original shapes but differing sizes, it might no longer be optimal.
+
+
 ## More details on paths
 
 Finding the optimal order of contraction is not an easy problem and formally scales factorially with respect to the number of terms in the expression. First, lets discuss what a path looks like in opt_einsum:
diff --git a/opt_einsum/__init__.py b/opt_einsum/__init__.py
index 2b3430f..16c844a 100644
--- a/opt_einsum/__init__.py
+++ b/opt_einsum/__init__.py
@@ -1,4 +1,4 @@
-from .contract import contract, contract_path
+from .contract import contract, contract_path, contract_expression
 from . import paths
 from . import blas
 from . import helpers
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 43b78e9..4efb1a9 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -259,7 +259,7 @@ def contract_path(*operands, **kwargs):
 def contract(*operands, **kwargs):
     """
     contract(subscripts, *operands, out=None, dtype=None, order='K',
-           casting='safe', use_blas=False, optimize=True, memory_limit=None)
+           casting='safe', use_blas=True, optimize=True, memory_limit=None)
 
     Evaluates the Einstein summation convention on the operands. A drop in
     replacment for NumPy's einsum function that optimizes the order of contraction
@@ -323,14 +323,10 @@ def contract(*operands, **kwargs):
     See opt_einsum.contract_path or numpy.einsum
 
     """
-
-    # Grab non-einsum kwargs
     optimize_arg = kwargs.pop('optimize', True)
     if optimize_arg is True:
         optimize_arg = 'greedy'
 
-    use_blas = kwargs.pop('use_blas', True)
-
     valid_einsum_kwargs = ['out', 'dtype', 'order', 'casting']
     einsum_kwargs = {k: v for (k, v) in kwargs.items() if k in valid_einsum_kwargs}
 
@@ -338,25 +334,38 @@ def contract(*operands, **kwargs):
     if optimize_arg is False:
         return np.einsum(*operands, **einsum_kwargs)
 
-    # Make sure all keywords are valid
-    valid_contract_kwargs = ['memory_limit', 'use_blas'] + valid_einsum_kwargs
-    unknown_kwargs = [k for (k, v) in kwargs.items() if k not in valid_contract_kwargs]
+    # Grab non-einsum kwargs
+    use_blas = kwargs.pop('use_blas', True)
+    memory_limit = kwargs.pop('memory_limit', None)
+    gen_expression = kwargs.pop('gen_expression', False)
+
+    # Make sure remaining keywords are valid for einsum
+    unknown_kwargs = [k for (k, v) in kwargs.items() if k not in valid_einsum_kwargs]
     if len(unknown_kwargs):
         raise TypeError("Did not understand the following kwargs: %s" % unknown_kwargs)
 
-    # Special handeling if out is specified
-    specified_out = False
-    out_array = einsum_kwargs.pop('out', None)
-    if out_array is not None:
-        specified_out = True
+    if gen_expression:
+        full_str = operands[0]
 
     # Build the contraction list and operand
-    memory_limit = kwargs.pop('memory_limit', None)
-
     operands, contraction_list = contract_path(
         *operands, path=optimize_arg, memory_limit=memory_limit, einsum_call=True, use_blas=use_blas)
 
-    handle_out = False
+    # check if performing contraction or just building expression
+    if gen_expression:
+        return ContractExpression(full_str, contraction_list, **einsum_kwargs)
+
+    return _core_contract(operands, contraction_list, **einsum_kwargs)
+
+
+def _core_contract(operands, contraction_list, **einsum_kwargs):
+    """Inner loop used to perform an actual contraction given the output
+    from a ``contract_path(..., einsum_call=True)`` call.
+    """
+
+    # Special handeling if out is specified
+    out_array = einsum_kwargs.pop('out', None)
+    specified_out = out_array is not None
 
     # Start contraction loop
     for num, contraction in enumerate(contraction_list):
@@ -366,8 +375,7 @@ def contract(*operands, **kwargs):
             tmp_operands.append(operands.pop(x))
 
         # Do we need to deal with the output?
-        if specified_out and ((num + 1) == len(contraction_list)):
-            handle_out = True
+        handle_out = specified_out and ((num + 1) == len(contraction_list))
 
         # Call tensordot
         if blas:
@@ -412,3 +420,104 @@ def contract(*operands, **kwargs):
         return out_array
     else:
         return operands[0]
+
+
+class ContractExpression:
+    """Helper class for storing an explicit ``contraction_list`` which can
+    then be repeatedly called solely with the array arguments.
+    """
+
+    def __init__(self, contraction, contraction_list, **einsum_kwargs):
+        self.contraction = contraction
+        self.contraction_list = contraction_list
+        self.einsum_kwargs = einsum_kwargs
+        self.num_args = len(contraction.split('->')[0].split(','))
+
+    def __call__(self, *arrays, **kwargs):
+        if len(arrays) != self.num_args:
+            raise ValueError("This `ContractExpression` takes exactly %s array arguments "
+                             "but received %s." % (self.num_args, len(arrays)))
+
+        out = kwargs.pop('out', None)
+        if kwargs:
+            raise ValueError("The only valid keyword argument to a `ContractExpression` "
+                             "call is `out=`. Got: %s." % kwargs)
+
+        try:
+            return _core_contract(list(arrays), self.contraction_list, out=out, **self.einsum_kwargs)
+        except ValueError as err:
+            original_msg = "".join(err.args) if err.args else ""
+            msg = ("Internal error while evaluating `ContractExpression`. Note that few checks are performed"
+                   " - the number and rank of the array arguments must match the original expression. "
+                   "The internal error was: '%s'" % original_msg, )
+            err.args = msg
+            raise
+
+    def __repr__(self):
+        return "ContractExpression('%s')" % self.contraction
+
+    def __str__(self):
+        s = "<ContractExpression> for '%s':" % self.contraction
+        for i, c in enumerate(self.contraction_list):
+            s += "\n  %i.  " % (i + 1)
+            s += "'%s'" % c[2] + (" [%s]" % c[-1] if c[-1] else "")
+        if self.einsum_kwargs:
+            s += "\neinsum_kwargs=%s" % self.einsum_kwargs
+        return s
+
+
+class _ShapeOnly(np.ndarray):
+    """Dummy ``numpy.ndarray`` which has a shape only - for generating
+    contract expressions.
+    """
+
+    def __init__(self, shape):
+        self.shape = shape
+
+
+def contract_expression(subscripts, *shapes, **kwargs):
+    """Generate an reusable expression for a given contraction with
+    specific shapes, which can for example be cached.
+
+    Parameters
+    ----------
+    subscripts : str
+        Specifies the subscripts for summation.
+    shapes : sequence of integer tuples
+        Shapes of the arrays to optimize the contraction for.
+    kwargs :
+        Passed on to ``contract_path`` or ``einsum``. See ``contract``.
+
+    Returns
+    -------
+    expr : ContractExpression
+        Callable with signature ``expr(*arrays)`` where the array's shapes
+        should match ``shapes``.
+
+    Notes
+    -----
+    - The `out` keyword argument should be supplied to the generated expression
+      rather than this function.
+    - The generated expression will work with any arrays which have
+      the same rank (number of dimensions) as the original shapes, however, if
+      the actual sizes are different, the expression may no longer be optimal.
+
+    Examples
+    --------
+
+    >>> expr = contract_expression("ab,bc->ac", (3, 4), (4, 5))
+    >>> a, b = np.random.rand(3, 4), np.random.rand(4, 5)
+    >>> c = expr(a, b)
+    >>> np.allclose(c, a @ b)
+    True
+
+    """
+    if not kwargs.get('optimize', True):
+        raise ValueError("Can only generate expressions for optimized contractions.")
+
+    if kwargs.get('out', None) is not None:
+        raise ValueError("`out` should only be specified when calling a `ContractExpression`, not when building it.")
+
+    dummy_arrays = [_ShapeOnly(s) for s in shapes]
+
+    return contract(subscripts, *dummy_arrays, gen_expression=True, **kwargs)


Test patch:
diff --git a/tests/test_contract.py b/tests/test_contract.py
index 016d784..6885522 100644
--- a/tests/test_contract.py
+++ b/tests/test_contract.py
@@ -5,7 +5,7 @@ Tets a series of opt_einsum contraction paths to ensure the results are the same
 from __future__ import division, absolute_import, print_function
 
 import numpy as np
-from opt_einsum import contract, contract_path, helpers
+from opt_einsum import contract, contract_path, helpers, contract_expression
 import pytest
 
 tests = [
@@ -127,3 +127,73 @@ def test_printing():
 
     ein = contract_path(string, *views)
     assert len(ein[1]) == 729
+
+
+@pytest.mark.parametrize("string", tests)
+@pytest.mark.parametrize("optimize", ['greedy', 'optimal'])
+@pytest.mark.parametrize("use_blas", [False, True])
+@pytest.mark.parametrize("out_spec", [False, True])
+def test_contract_expressions(string, optimize, use_blas, out_spec):
+    views = helpers.build_views(string)
+    shapes = [view.shape for view in views]
+    expected = contract(string, *views, optimize=False, use_blas=False)
+
+    expr = contract_expression(
+        string, *shapes, optimize=optimize, use_blas=use_blas)
+
+    if out_spec and ("->" in string) and (string[-2:] != "->"):
+        out, = helpers.build_views(string.split('->')[1])
+        expr(*views, out=out)
+    else:
+        out = expr(*views)
+
+    assert np.allclose(out, expected)
+
+    # check representations
+    assert string in expr.__repr__()
+    assert string in expr.__str__()
+
+
+def test_contract_expression_checks():
+    # check optimize needed
+    with pytest.raises(ValueError):
+        contract_expression("ab,bc->ac", (2, 3), (3, 4), optimize=False)
+
+    # check sizes are still checked
+    with pytest.raises(ValueError):
+        contract_expression("ab,bc->ac", (2, 3), (3, 4), (42, 42))
+
+    # check if out given
+    out = np.empty((2, 4))
+    with pytest.raises(ValueError):
+        contract_expression("ab,bc->ac", (2, 3), (3, 4), out=out)
+
+    # check still get errors when wrong ranks supplied to expression
+    expr = contract_expression("ab,bc->ac", (2, 3), (3, 4))
+
+    # too few arguments
+    with pytest.raises(ValueError) as err:
+        expr(np.random.rand(2, 3))
+    assert "`ContractExpression` takes exactly 2" in str(err)
+
+    # too many arguments
+    with pytest.raises(ValueError) as err:
+        expr(np.random.rand(2, 3), np.random.rand(2, 3), np.random.rand(2, 3))
+    assert "`ContractExpression` takes exactly 2" in str(err)
+
+    # wrong shapes
+    with pytest.raises(ValueError) as err:
+        expr(np.random.rand(2, 3, 4), np.random.rand(3, 4))
+    assert "Internal error while evaluating `ContractExpression`" in str(err)
+    with pytest.raises(ValueError) as err:
+        expr(np.random.rand(2, 4), np.random.rand(3, 4, 5))
+    assert "Internal error while evaluating `ContractExpression`" in str(err)
+    with pytest.raises(ValueError) as err:
+        expr(np.random.rand(2, 3), np.random.rand(3, 4),
+             out=np.random.rand(2, 4, 6))
+    assert "Internal error while evaluating `ContractExpression`" in str(err)
+
+    # should only be able to specify out
+    with pytest.raises(ValueError) as err:
+        expr(np.random.rand(2, 3), np.random.rand(3, 4), order='F')
+    assert "only valid keyword argument to a `ContractExpression`" in str(err)