Repository: dgasmith/opt_einsum (python)

Issue:
Re-use intermediates across various contractions (Common Subexpression Elimination)
Suppose you want to compute two contractions of the same tensors, e.g. contraction strings

`['ab,dca,eb,cde', 'ab,cda,eb,cde']`

The (globally) optimal way to do this would be to first perform the contractions over indices a,b and e, and then perform the remaining contractions over c and d for the two sets of contractions. The current opt_einsum implementation does not allow for such a _global_ optimization of contraction order and re-use of common intermediates.

I'm still organizing my thoughts on this, and all input would be most welcome. On a side note, I believe implementing such a more general optimization strategy will also fix #7 as a by-product.

### Optimization logic

A relevant publication suggesting a concrete algorithm for this optimization problem is

[_Hartono et al_, Identifying Cost-Effective Common Subexpressions to Reduce Operation Count in Tensor Contraction Evaluations](http://www.csc.lsu.edu/~gb/TCE/Publications/OpMinCSE-ICCS06.pdf)

I do not know to what extent the current code can be re-used in this more general setup, but the _single-term_ optimization should be under control with the current implementation.

### Interface

Such a multi_contract function could be called with a list of tuples (contractionString, [tensors, ..]), and would return a list of results with the same length as the input list.

Internally, one would have to find out which tensors are actually identical between the various contractions, and then use the above contraction logic. Ideally this information should be deduced automatically and not rely on user input being in the correct form. In the same spirit, dummy indices should be allowed to have arbitrary names, i.e. they should not have to match across different contractions to be correctly identified as a common subexpression.
This may require transforming the input into a 'canonical' form first to make sure that common subexpressions are captured correctly.

In contrast to the setup outlined in _Hartono et al_, contraction strings should maybe remain in their current 'simple' form and **not** be generalized to allow for numerical expressions like sums of tensors etc. Such a behavior can be implemented _a posteriori_ with the interface described here by computing the sum of the resulting tensors, e.g.

`contract('ab,ab + ab,ba', N,M) -->  'sum'(multi_contract( [('ab,ab', N,M), ('ab,ba', N,M)] ))`

Thus, restricting contraction strings to be of the form currently used does not cause loss of generality, the only downside being that it might lead to a slightly increased memory-footprint as the function would return several arrays instead of one.


Other thoughts?

Pull request:
Implement shared_intermediates context manager
Resolves #9 

## Description

This implements a `shared_intermediates` context manager, within which multiple calls to `contract` will share intermediate computations. This is very useful e.g. in message passing algorithms. The first draft of this PR is copied directly [from Pyro](https://github.com/uber/pyro/blob/1b1910eadff80178197879bf847beb4443fd9014/pyro/ops/einsum/shared.py), where we are using sharing for probabilistic inference in graphical models.

The implementation uses a special internal backend `opt_einsum.backends.shared`. When a sharing context is opened
```py
with shared_intermediates():
    contract(..., backend='foo')
    contract(..., backend='foo')
```
a special `handle_sharing()` context is activated inside the `contract` function, temporarily switching to the `shared` backend, and setting the original backend `foo` in a global `shared._CURRENT_BACKEND` variable. The sharing backend then performs all the original `foo` operations and also memoizes them.

One design choice is to memoize by `id()` rather than by value. This makes sense from a computational perspective (equality comparison is expensive), but requires a bit more care by users.

A second design choice is to expose the memoization `cache` to users. This makes it easy to share intermediates by passing the cache around, as is necessary when it is difficult to perform all computation in a context:
```py
with shared_intermediates() as cache:
    contract(...)
...pass control elsewhere...
with shared_intermediates(cache):  # <-- reuse previous cache
    contract(...)
del cache  # <-- intermediates will be garbage collected
```

## Todos
- [x] Cache type conversions e.g. `numpy`->`torch`->`numpy` so args can be compared by id.
- [x] fix failing tests
- [x] add tests that pass around the cache
- [x] add tests where we tests the normal ordering:
    ```py
    with shared_intermeidates() as cache:
        a = np.random.rand(4, 4)
        b = np.random.rand(4, 4)
        contract("ab,bc->ac", a, b)

    assert get_cache_key("zc,az->ac", b, a) in cache
    ...
    ```
- [x] test against the Pyro [rm-einsum-shared](https://github.com/uber/pyro/pull/1327) branch
- [x] test nested sharing
- [x] write docs
- [x] move `_alpha_canonicalize()` to `parser.py`
- [x] test caching with constant expressions

## Questions
- [x] Could a maintainer suggest where the `handle_sharing()` context should be moved? I've attempted to insert it into `ContractExpression.__call__()`, but I'm a bit lost.

## Status
- [x] Ready to go

Patch:
diff --git a/README.md b/README.md
index 8c3d13b..69e3774 100644
--- a/README.md
+++ b/README.md
@@ -19,6 +19,7 @@ As well as [`opt_einsum.contract`](https://optimized-einsum.readthedocs.io/en/la
 * Perform the contractions with many [different backends](http://optimized-einsum.readthedocs.io/en/latest/backends.html), including on the GPU and with libraries such as [TensorFlow](https://www.tensorflow.org) and [PyTorch](https://pytorch.org).
 * Generate [reusable expressions](http://optimized-einsum.readthedocs.io/en/latest/reusing_paths.html), potentially with [constant tensors](http://optimized-einsum.readthedocs.io/en/latest/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](http://optimized-einsum.readthedocs.io/en/latest/ex_large_expr_with_greedy.html).
+* Share [intermediate computations](http://optimized-einsum.readthedocs.io/en/latest/sharing_intermediates.html) among multiple contractions.
 
 ## Quick tutorial
 Einsum is a powerful function for contracting tensors of arbitrary dimension and index.
diff --git a/docs/source/api.rst b/docs/source/api.rst
index 1a293c0..dd59a24 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -12,3 +12,4 @@ Function Reference
     opt_einsum.paths.optimal
     opt_einsum.paths.greedy
     opt_einsum.parser.get_symbol
+    opt_einsum.sharing.shared_intermediates
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 5948454..869a28f 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -123,7 +123,7 @@ Table of Contents
    install
    backends
    reusing_paths
-
+   sharing_intermediates
 
 .. toctree::
    :maxdepth: 1
diff --git a/docs/source/sharing_intermediates.rst b/docs/source/sharing_intermediates.rst
new file mode 100644
index 0000000..345a2c4
--- /dev/null
+++ b/docs/source/sharing_intermediates.rst
@@ -0,0 +1,41 @@
+=====================
+Sharing Intermediates
+=====================
+
+If you want to compute multiple similar contractions with common terms, you can embed them in a :func:`~opt_einsum.shared_intermediates` context. Computations of subexpressions in this context will be memoized, and will be garbage collected when the contexts exits.
+
+For example, suppose we want to compute marginals at each point in a factor chain:
+
+.. code:: python
+
+    inputs = 'ab,bc,cd,de,ef'
+    factors = [np.random.rand(1000, 1000) for _ in range(5)]
+
+    %%timeit
+    marginals = {output: contract('{}->{}'.format(inputs, output), *factors)
+                 for output in 'abcdef'}
+    1 loop, best of 3: 5.82 s per loop
+
+To share this computation, we can perform all contractions in a shared context
+
+.. code:: python
+
+    %%timeit
+    with shared_intermediates():
+        marginals = {output: contract('{}->{}'.format(inputs, output), *factors)
+                     for output in 'abcdef'}
+    1 loop, best of 3: 1.55 s per loop
+
+If it is difficult to fit your code into a context, you can instead save the sharing cache for later reuse.
+
+.. code:: python
+
+    with shared_intermediates() as cache:  # create a cache
+        pass
+    marginals = {}
+    for output in 'abcdef':
+        with shared_intermediates(cache):  # reuse a common cache
+            marginals[output] = contract('{}->{}'.format(inputs, output), *factors)
+    del cache  # garbage collect intermediates
+
+Note that sharing contexts can be nested, so it is safe to to use :func:`~opt_einsum.shared_intermediates` in library code without leaking intermediates into user caches.
diff --git a/opt_einsum/__init__.py b/opt_einsum/__init__.py
index 8367dc8..1db2b64 100644
--- a/opt_einsum/__init__.py
+++ b/opt_einsum/__init__.py
@@ -4,6 +4,7 @@ Main init function for opt_einsum.
 
 from .contract import contract, contract_path, contract_expression
 from .parser import get_symbol
+from .sharing import shared_intermediates
 from . import paths
 from . import blas
 from . import helpers
diff --git a/opt_einsum/backends/cupy.py b/opt_einsum/backends/cupy.py
index 1fa4a4f..ce7bf6b 100644
--- a/opt_einsum/backends/cupy.py
+++ b/opt_einsum/backends/cupy.py
@@ -4,8 +4,10 @@ Required functions for optimized contractions of numpy arrays using cupy.
 
 from __future__ import absolute_import
 import numpy as np
+from ..sharing import to_backend_cache_wrap
 
 
+@to_backend_cache_wrap
 def to_cupy(array):  # pragma: no cover
     import cupy
 
diff --git a/opt_einsum/backends/torch.py b/opt_einsum/backends/torch.py
index 5d893d4..124d87b 100644
--- a/opt_einsum/backends/torch.py
+++ b/opt_einsum/backends/torch.py
@@ -6,6 +6,7 @@ from __future__ import absolute_import
 import numpy as np
 
 from ..parser import convert_to_valid_einsum_chars, einsum_symbols_base
+from ..sharing import to_backend_cache_wrap
 
 _TORCH_DEVICE = None
 
@@ -84,6 +85,7 @@ def tensordot(x, y, axes=2):
     return einsum(einsum_str, x, y)
 
 
+@to_backend_cache_wrap
 def to_torch(array):
     torch, device = _get_torch_and_device()
 
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 9971bf1..a7932f0 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -9,6 +9,7 @@ from . import blas
 from . import helpers
 from . import parser
 from . import paths
+from . import sharing
 
 
 def contract_path(*operands, **kwargs):
@@ -267,6 +268,7 @@ def contract_path(*operands, **kwargs):
     return path, path_print
 
 
+@sharing.einsum_cache_wrap
 def _einsum(*operands, **kwargs):
     """Base einsum, but with pre-parse for valid characters if string given.
     """
@@ -289,6 +291,7 @@ def _einsum(*operands, **kwargs):
     return fn(einsum_str, *operands, **kwargs)
 
 
+@sharing.transpose_cache_wrap
 def _transpose(x, axes, backend='numpy'):
     """Base transpose.
     """
@@ -300,6 +303,7 @@ def _transpose(x, axes, backend='numpy'):
         return fn(x, axes)
 
 
+@sharing.tensordot_cache_wrap
 def _tensordot(x, y, axes, backend='numpy'):
     """Base tensordot.
     """
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index a684089..c5370b1 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -4,8 +4,9 @@
 A functionally equivalent parser of the numpy.einsum input parser
 """
 
-import numpy as np
+from collections import OrderedDict
 
+import numpy as np
 
 einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 
@@ -65,6 +66,18 @@ def convert_to_valid_einsum_chars(einsum_str):
     return "".join(replacer.get(x, x) for x in einsum_str)
 
 
+def alpha_canonicalize(equation):
+    """Alpha convert an equation in an order-independent canonical way.
+    """
+    rename = OrderedDict()
+    for name in equation:
+        if name in '.,->':
+            continue
+        if name not in rename:
+            rename[name] = get_symbol(len(rename))
+    return ''.join(rename.get(x, x) for x in equation)
+
+
 def find_output_str(subscripts):
     """Find the output string for the inputs ``susbcripts``.
     """
diff --git a/opt_einsum/sharing.py b/opt_einsum/sharing.py
new file mode 100644
index 0000000..504ba05
--- /dev/null
+++ b/opt_einsum/sharing.py
@@ -0,0 +1,153 @@
+# Copyright (c) 2018 Uber Technologies
+
+import contextlib
+import functools
+import numbers
+from collections import Counter
+
+from .parser import alpha_canonicalize, parse_einsum_input
+
+_SHARING_STACK = []
+
+
+@contextlib.contextmanager
+def shared_intermediates(cache=None):
+    """Context in which contract intermediate results are shared.
+
+    Note that intermediate computations will not be garbage collected until
+    1. this context exits, and
+    2. the yielded cache is garbage collected (if it was captured).
+
+    Parameters
+    ----------
+    cache : dict
+        If specified, a user-stored dict in which intermediate results will
+        be stored. This can be used to interleave sharing contexts.
+
+    Returns
+    -------
+    cache : dict
+        A dictionary in which sharing results are stored. If ignored,
+        sharing results will be garbage collected when this context is
+        exited. This dict can be passed to another context to resume
+        sharing.
+    """
+    if cache is None:
+        cache = {}
+    try:
+        _SHARING_STACK.append(cache)
+        yield cache
+    finally:
+        _SHARING_STACK.pop()
+
+
+def count_cached_ops(cache):
+    """Returns a counter of the types of each op in the cache.
+    This is useful for profiling to increase sharing.
+    """
+    return Counter(key[0] for key in cache.keys())
+
+
+def _save_tensors(*tensors):
+    """Save tensors in the cache to prevent their ids from being recycled.
+    This is needed to prevent false cache lookups.
+    """
+    cache = _SHARING_STACK[-1]
+    for tensor in tensors:
+        cache['tensor', id(tensor)] = tensor
+
+
+def _memoize(key, fn, *args, **kwargs):
+    """Memoize ``fn(*args, **kwargs)`` using the given ``key``.
+    Results will be stored in the innermost ``cache`` yielded by
+    :func:`shared_intermediates`.
+    """
+    cache = _SHARING_STACK[-1]
+    if key in cache:
+        return cache[key]
+    result = fn(*args, **kwargs)
+    cache[key] = result
+    return result
+
+
+def transpose_cache_wrap(transpose):
+    """Decorates a ``transpose()`` implementation to be memoized inside a
+    :func:`shared_intermediates` context.
+    """
+
+    @functools.wraps(transpose)
+    def cached_transpose(a, axes, backend='numpy'):
+        if not _SHARING_STACK:
+            return transpose(a, axes, backend=backend)
+
+        # hash by axes
+        _save_tensors(a)
+        axes = tuple(axes)
+        key = 'transpose', backend, id(a), axes
+        return _memoize(key, transpose, a, axes, backend=backend)
+
+    return cached_transpose
+
+
+def tensordot_cache_wrap(tensordot):
+    """Decorates a ``tensordot()`` implementation to be memoized inside a
+    :func:`shared_intermediates` context.
+    """
+
+    @functools.wraps(tensordot)
+    def cached_tensordot(x, y, axes=2, backend='numpy'):
+        if not _SHARING_STACK:
+            return tensordot(x, y, axes, backend=backend)
+
+        # hash based on the (axes_x,axes_y) form of axes
+        _save_tensors(x, y)
+        if isinstance(axes, numbers.Number):
+            axes = list(range(len(x.shape)))[len(x.shape) - axes:], list(range(len(y.shape)))[:axes]
+        axes = tuple(axes[0]), tuple(axes[1])
+        key = 'tensordot', backend, id(x), id(y), axes
+        return _memoize(key, tensordot, x, y, axes, backend=backend)
+
+    return cached_tensordot
+
+
+def einsum_cache_wrap(einsum):
+    """Decorates an ``einsum()`` implementation to be memoized inside a
+    :func:`shared_intermediates` context.
+    """
+
+    @functools.wraps(einsum)
+    def cached_einsum(*args, **kwargs):
+        if not _SHARING_STACK:
+            return einsum(*args, **kwargs)
+
+        # hash modulo commutativity by computing a canonical ordering and names
+        backend = kwargs.pop('backend', 'numpy')
+        equation = args[0]
+        inputs, output, operands = parse_einsum_input(args)
+        inputs = inputs.split(',')
+        _save_tensors(*operands)
+        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))
+        key = 'einsum', backend, canonical_equation, canonical_ids
+        return _memoize(key, einsum, equation, *operands, backend=backend)
+
+    return cached_einsum
+
+
+def to_backend_cache_wrap(to_backend):
+    """Decorates an ``to_backend()`` implementation to be memoized inside a
+    :func:`shared_intermediates` context (e.g. ``to_cupy``, ``to_torch``).
+    """
+
+    @functools.wraps(to_backend)
+    def cached_to_backend(array):
+        if not _SHARING_STACK:
+            return to_backend(array)
+
+        # hash by id
+        key = to_backend.__name__, id(array)
+        return _memoize(key, to_backend, array)
+
+    return cached_to_backend


Test patch:
diff --git a/opt_einsum/tests/test_sharing.py b/opt_einsum/tests/test_sharing.py
new file mode 100644
index 0000000..b0b860a
--- /dev/null
+++ b/opt_einsum/tests/test_sharing.py
@@ -0,0 +1,364 @@
+import itertools
+import weakref
+from collections import Counter
+
+import numpy as np
+import pytest
+
+from opt_einsum import (contract_expression, contract_path, get_symbol,
+                        helpers, shared_intermediates)
+from opt_einsum.backends import to_cupy, to_torch
+from opt_einsum.contract import _einsum
+from opt_einsum.parser import parse_einsum_input
+from opt_einsum.sharing import count_cached_ops
+
+try:
+    import cupy
+    cupy_if_found = 'cupy'
+except ImportError:
+    cupy_if_found = pytest.param('cupy', marks=[pytest.mark.skip(reason="CuPy not installed.")])
+
+try:
+    import torch
+    torch_if_found = 'torch'
+except ImportError:
+    torch_if_found = pytest.param('torch', marks=[pytest.mark.skip(reason="PyTorch not installed.")])
+
+backends = ['numpy', torch_if_found, cupy_if_found]
+equations = [
+    'ab,bc->ca',
+    'abc,bcd,dea',
+    'abc,def->fedcba',
+    'abc,bcd,df->fa',
+    # test 'prefer einsum' ops
+    'ijk,ikj',
+    'i,j->ij',
+    'ijk,k->ij',
+    'AB,BC->CA',
+]
+to_backend = {
+    'numpy': lambda x: x,
+    'torch': to_torch,
+    'cupy': to_cupy,
+}
+
+
+@pytest.mark.parametrize('eq', equations)
+@pytest.mark.parametrize('backend', backends)
+def test_sharing_value(eq, backend):
+    views = helpers.build_views(eq)
+    shapes = [v.shape for v in views]
+    expr = contract_expression(eq, *shapes)
+
+    expected = expr(*views, backend=backend)
+    with shared_intermediates():
+        actual = expr(*views, backend=backend)
+
+    assert (actual == expected).all()
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_complete_sharing(backend):
+    eq = 'ab,bc,cd->'
+    views = helpers.build_views(eq)
+    expr = contract_expression(eq, *(v.shape for v in views))
+
+    print('-' * 40)
+    print('Without sharing:')
+    with shared_intermediates() as cache:
+        expr(*views, backend=backend)
+        expected = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('With sharing:')
+    with shared_intermediates() as cache:
+        expr(*views, backend=backend)
+        expr(*views, backend=backend)
+        actual = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('Without sharing: {} expressions'.format(expected))
+    print('With sharing: {} expressions'.format(actual))
+    assert actual == expected
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_sharing_reused_cache(backend):
+    eq = 'ab,bc,cd->'
+    views = helpers.build_views(eq)
+    expr = contract_expression(eq, *(v.shape for v in views))
+
+    print('-' * 40)
+    print('Without sharing:')
+    with shared_intermediates() as cache:
+        expr(*views, backend=backend)
+        expected = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('With sharing:')
+    with shared_intermediates() as cache:
+        expr(*views, backend=backend)
+    with shared_intermediates(cache):
+        expr(*views, backend=backend)
+        actual = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('Without sharing: {} expressions'.format(expected))
+    print('With sharing: {} expressions'.format(actual))
+    assert actual == expected
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_no_sharing_separate_cache(backend):
+    eq = 'ab,bc,cd->'
+    views = helpers.build_views(eq)
+    expr = contract_expression(eq, *(v.shape for v in views))
+
+    print('-' * 40)
+    print('Without sharing:')
+    with shared_intermediates() as cache:
+        expr(*views, backend=backend)
+        expected = count_cached_ops(cache)
+        expected.update(count_cached_ops(cache))  # we expect double
+
+    print('-' * 40)
+    print('With sharing:')
+    with shared_intermediates() as cache1:
+        expr(*views, backend=backend)
+        actual = count_cached_ops(cache1)
+    with shared_intermediates() as cache2:
+        expr(*views, backend=backend)
+        actual.update(count_cached_ops(cache2))
+
+    print('-' * 40)
+    print('Without sharing: {} expressions'.format(expected))
+    print('With sharing: {} expressions'.format(actual))
+    assert actual == expected
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_sharing_nesting(backend):
+    eqs = ['ab,bc,cd->a',
+           'ab,bc,cd->b',
+           'ab,bc,cd->c',
+           'ab,bc,cd->c']
+    views = helpers.build_views(eqs[0])
+    shapes = [v.shape for v in views]
+    refs = weakref.WeakValueDictionary()
+
+    def method1(views):
+        with shared_intermediates():
+            w = contract_expression(eqs[0], *shapes)(*views, backend=backend)
+            x = contract_expression(eqs[2], *shapes)(*views, backend=backend)
+            result = contract_expression('a,b->', w.shape, x.shape)(w, x, backend=backend)
+            refs['w'] = w
+            refs['x'] = x
+            del w, x
+            assert 'w' in refs
+            assert 'x' in refs
+        assert 'w' not in refs, 'cache leakage'
+        assert 'x' not in refs, 'cache leakage'
+        return result
+
+    def method2(views):
+        with shared_intermediates():
+            y = contract_expression(eqs[2], *shapes)(*views, backend=backend)
+            z = contract_expression(eqs[3], *shapes)(*views, backend=backend)
+            refs['y'] = y
+            refs['z'] = z
+            result = contract_expression('c,d->', y.shape, z.shape)(y, z, backend=backend)
+            result = result + method1(views)  # nest method1 in method2
+            del y, z
+            assert 'y' in refs
+            assert 'z' in refs
+        assert 'y' not in refs
+        assert 'z' not in refs
+
+    method1(views)
+    method2(views)
+
+
+@pytest.mark.parametrize('eq', equations)
+@pytest.mark.parametrize('backend', backends)
+def test_sharing_modulo_commutativity(eq, backend):
+    ops = helpers.build_views(eq)
+    ops = [to_backend[backend](x) for x in ops]
+    inputs, output, _ = parse_einsum_input([eq] + ops)
+    inputs = inputs.split(',')
+
+    print('-' * 40)
+    print('Without sharing:')
+    with shared_intermediates() as cache:
+        _einsum(eq, *ops, backend=backend)
+        expected = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('With sharing:')
+    with shared_intermediates() as cache:
+        for permuted in itertools.permutations(zip(inputs, ops)):
+            permuted_inputs = [p[0] for p in permuted]
+            permuted_ops = [p[1] for p in permuted]
+            permuted_eq = '{}->{}'.format(','.join(permuted_inputs), output)
+            _einsum(permuted_eq, *permuted_ops, backend=backend)
+        actual = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('Without sharing: {} expressions'.format(expected))
+    print('With sharing: {} expressions'.format(actual))
+    assert actual == expected
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_partial_sharing(backend):
+    eq = 'ab,bc,de->'
+    x, y, z1 = helpers.build_views(eq)
+    z2 = 2.0 * z1 - 1.0
+    expr = contract_expression(eq, x.shape, y.shape, z1.shape)
+
+    print('-' * 40)
+    print('Without sharing:')
+    num_exprs_nosharing = Counter()
+    with shared_intermediates() as cache:
+        expr(x, y, z1, backend=backend)
+        num_exprs_nosharing.update(count_cached_ops(cache))
+    with shared_intermediates() as cache:
+        expr(x, y, z2, backend=backend)
+        num_exprs_nosharing.update(count_cached_ops(cache))
+
+    print('-' * 40)
+    print('With sharing:')
+    with shared_intermediates() as cache:
+        expr(x, y, z1, backend=backend)
+        expr(x, y, z2, backend=backend)
+        num_exprs_sharing = count_cached_ops(cache)
+
+    print('-' * 40)
+    print('Without sharing: {} expressions'.format(num_exprs_nosharing))
+    print('With sharing: {} expressions'.format(num_exprs_sharing))
+    assert num_exprs_nosharing['einsum'] > num_exprs_sharing['einsum']
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_sharing_with_constants(backend):
+    inputs = 'ij,jk,kl'
+    outputs = 'ijkl'
+    equations = ['{}->{}'.format(inputs, output) for output in outputs]
+    shapes = (2, 3), (3, 4), (4, 5)
+    constants = {0, 2}
+    ops = [np.random.rand(*shp) if i in constants else shp for i, shp in enumerate(shapes)]
+    var = np.random.rand(*shapes[1])
+
+    expected = [contract_expression(eq, *shapes)(ops[0], var, ops[2])
+                for eq in equations]
+
+    with shared_intermediates():
+        actual = [contract_expression(eq, *ops, constants=constants)(var)
+                  for eq in equations]
+
+    for dim, expected_dim, actual_dim in zip(outputs, expected, actual):
+        assert np.allclose(expected_dim, actual_dim), 'error at {}'.format(dim)
+
+
+@pytest.mark.parametrize('size', [3, 4, 5])
+@pytest.mark.parametrize('backend', backends)
+def test_chain(size, backend):
+    xs = [np.random.rand(2, 2) for _ in range(size)]
+    shapes = [x.shape for x in xs]
+    alphabet = ''.join(get_symbol(i) for i in range(size + 1))
+    names = [alphabet[i:i+2] for i in range(size)]
+    inputs = ','.join(names)
+
+    with shared_intermediates():
+        print(inputs)
+        for i in range(size + 1):
+            target = alphabet[i]
+            eq = '{}->{}'.format(inputs, target)
+            path_info = contract_path(eq, *xs)
+            print(path_info[1])
+            expr = contract_expression(eq, *shapes)
+            expr(*xs, backend=backend)
+        print('-' * 40)
+
+
+@pytest.mark.parametrize('size', [3, 4, 5, 10])
+@pytest.mark.parametrize('backend', backends)
+def test_chain_2(size, backend):
+    xs = [np.random.rand(2, 2) for _ in range(size)]
+    shapes = [x.shape for x in xs]
+    alphabet = ''.join(get_symbol(i) for i in range(size + 1))
+    names = [alphabet[i:i+2] for i in range(size)]
+    inputs = ','.join(names)
+
+    with shared_intermediates():
+        print(inputs)
+        for i in range(size):
+            target = alphabet[i:i+2]
+            eq = '{}->{}'.format(inputs, target)
+            path_info = contract_path(eq, *xs)
+            print(path_info[1])
+            expr = contract_expression(eq, *shapes)
+            expr(*xs, backend=backend)
+        print('-' * 40)
+
+
+def _compute_cost(cache):
+    counts = count_cached_ops(cache)
+    return counts['einsum'] + counts['tensordot']
+
+
+@pytest.mark.parametrize('backend', backends)
+def test_chain_2_growth(backend):
+    sizes = list(range(1, 21))
+    costs = []
+    for size in sizes:
+        xs = [np.random.rand(2, 2) for _ in range(size)]
+        alphabet = ''.join(get_symbol(i) for i in range(size + 1))
+        names = [alphabet[i:i+2] for i in range(size)]
+        inputs = ','.join(names)
+
+        with shared_intermediates() as cache:
+            for i in range(size):
+                target = alphabet[i:i+2]
+                eq = '{}->{}'.format(inputs, target)
+                expr = contract_expression(eq, *(x.shape for x in xs))
+                expr(*xs, backend=backend)
+            costs.append(_compute_cost(cache))
+
+    print('sizes = {}'.format(repr(sizes)))
+    print('costs = {}'.format(repr(costs)))
+    for size, cost in zip(sizes, costs):
+        print('{}\t{}'.format(size, cost))
+
+
+@pytest.mark.parametrize('size', [3, 4, 5])
+@pytest.mark.parametrize('backend', backends)
+def test_chain_sharing(size, backend):
+    xs = [np.random.rand(2, 2) for _ in range(size)]
+    alphabet = ''.join(get_symbol(i) for i in range(size + 1))
+    names = [alphabet[i:i+2] for i in range(size)]
+    inputs = ','.join(names)
+
+    num_exprs_nosharing = 0
+    for i in range(size + 1):
+        with shared_intermediates() as cache:
+            target = alphabet[i]
+            eq = '{}->{}'.format(inputs, target)
+            expr = contract_expression(eq, *(x.shape for x in xs))
+            expr(*xs, backend=backend)
+            num_exprs_nosharing += _compute_cost(cache)
+
+    with shared_intermediates() as cache:
+        print(inputs)
+        for i in range(size + 1):
+            target = alphabet[i]
+            eq = '{}->{}'.format(inputs, target)
+            path_info = contract_path(eq, *xs)
+            print(path_info[1])
+            expr = contract_expression(eq, *(x.shape for x in xs))
+            expr(*xs, backend=backend)
+        num_exprs_sharing = _compute_cost(cache)
+
+    print('-' * 40)
+    print('Without sharing: {} expressions'.format(num_exprs_nosharing))
+    print('With sharing: {} expressions'.format(num_exprs_sharing))
+    assert num_exprs_nosharing > num_exprs_sharing