Repository: dgasmith/opt_einsum
Pull request #60 (2018-09-30 18:41:43)

Add prototype cheap optimizer
Addresses #56 
Mirrors ~~https://github.com/uber/pyro/pull/1424~~ https://github.com/uber/pyro/pull/1430

## Description

Implements a cheap approximate optimizer for very large paths. This is currently over 100x faster on some large contractions in Pyro, but produces slightly suboptimal paths.

## Todos
- [x] Fix issues with determinism, required by `shared_intermediates`
- [x] Ensure algorithm is not much worse than `greedy`
- ~~Add quality tests (will be easier after #61)~~

## Questions
- [x] How should we be computing the final outer products? Can we do this all in one go?
- [x] What tests should I add?

## Status
- [x] Ready to go

codecov-io: # [Codecov](https://codecov.io/gh/dgasmith/opt_einsum/pull/60?src=pr&el=h1) Report
> Merging [#60](https://codecov.io/gh/dgasmith/opt_einsum/pull/60?src=pr&el=desc) into [master](https://codecov.io/gh/dgasmith/opt_einsum/commit/1db5a0ab3dd8003919d0fa38b455cc7bad1b4e80?src=pr&el=desc) will **decrease** coverage by `0.6%`.
> The diff coverage is `91.07%`.


jcmgray: I've updated a gist I have of contraction finding time [here](https://gist.github.com/jcmgray/2fb88f1b8473eef5070316f02e26d66e), it looks good for large contractions!

![index](https://user-images.githubusercontent.com/8982598/46289199-417d0c00-c580-11e8-9153-fcfdd6e33324.png).

In a 'real-world' contraction, there is a slight difference in timings, maybe due to ignoring FLOP cost? Might also just be random though.
jcmgray: Do we want this to replace the current ``'greedy'`` method? I think it probably could, and as long as it maintains the ``size12 - size1 - size2`` cost it should find essentially the same contraction paths (modulo a bit of variance from choosing different branches when several have the same cost).

The one remaining difference (I think?) is that the current 'greedy' breaks ties using the flop-cost, but adding this back in might make most sense using a custom cost function like proposed in #64.

We could keep the current greedy around as ``'greedy_legacy'`` or something just for reference.
fritzo: Ok, I've addressed review comments. The current cost function is `size12 - (size1 + size2) / 2`, since this seems to work on the few sharing examples I've tried. I don't mind if you change it back to `size12 - size1 - size2`.

I think it's reasonable to merge (pending further comments) and then build some experimentation tooling to determine what objective function works will. One thing that would be handy would be a FLOP counter that worked with sharing; currently I'm running an actual computation and counting number of things in the sharing `cache`.
jcmgray: All looks great to me. I've updated this gist of [timings and other tests here](https://gist.github.com/jcmgray/2fb88f1b8473eef5070316f02e26d66e) in case its of interest (main timings similar to above still).

Patch:
diff --git a/docs/source/api.rst b/docs/source/api.rst
index dd59a24..935e907 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -11,5 +11,6 @@ Function Reference
     opt_einsum.contract.ContractExpression
     opt_einsum.paths.optimal
     opt_einsum.paths.greedy
+    opt_einsum.paths.eager
     opt_einsum.parser.get_symbol
     opt_einsum.sharing.shared_intermediates
diff --git a/docs/source/eager_path.rst b/docs/source/eager_path.rst
new file mode 100644
index 0000000..8b4b25f
--- /dev/null
+++ b/docs/source/eager_path.rst
@@ -0,0 +1,13 @@
+==============
+The Eager Path
+==============
+
+The ``eager`` path is constructed in three stages:
+
+1. Eagerly compute Hadamard products (in arbitrary order -- this is commutative).
+2. Greedily contract pairs of remaining tensors, at each step choosing the pair that maximizes ``reduced_size``.
+3. Greedily compute pairwise outer products, at each step choosing the pair that minimizes ``sum(input_sizes)``.
+
+The ``eager`` algorithm has space and time complexity ``O(n * k)`` where ``n`` is the number of input tensors and ``k`` is the maximum number of tensors that share any dimension (excluding dimensions that occurr in the output or in every tensor).
+This algorithm this scales well to very large sparse contractions of small tensors.
+The ``eager`` functionality is provided by :func:`~opt_einsum.paths.eager`.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 400efbe..7b8b92b 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -134,6 +134,7 @@ Table of Contents
    path_finding
    optimal_path
    greedy_path
+   eager_path
 
 .. toctree::
    :maxdepth: 1
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 7e5e233..4a42176 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -42,23 +42,23 @@ class PathInfo(object):
         header = ("scaling", "BLAS", "current", "remaining")
 
         path_print = [
-            "  Complete contraction:  {}\n".format(overall_contraction),
-            "         Naive scaling:  {}\n".format(len(self.indices)),
-            "     Optimized scaling:  {}\n".format(max(self.scale_list)),
-            "      Naive FLOP count:  {:.3e}\n".format(naive_cost),
-            "  Optimized FLOP count:  {:.3e}\n".format(opt_cost),
-            "   Theoretical speedup:  {:3.3f}\n".format(speedup),
-            "  Largest intermediate:  {:.3e} elements\n".format(largest_intermediate),
-            "-" * 80 + "\n",
-            "{:>6} {:>11} {:>22} {:>37}\n".format(*header),
-            "-" * 80
+            u"  Complete contraction:  {}\n".format(overall_contraction),
+            u"         Naive scaling:  {}\n".format(len(self.indices)),
+            u"     Optimized scaling:  {}\n".format(max(self.scale_list)),
+            u"      Naive FLOP count:  {:.3e}\n".format(naive_cost),
+            u"  Optimized FLOP count:  {:.3e}\n".format(opt_cost),
+            u"   Theoretical speedup:  {:3.3f}\n".format(speedup),
+            u"  Largest intermediate:  {:.3e} elements\n".format(largest_intermediate),
+            u"-" * 80 + "\n",
+            u"{:>6} {:>11} {:>22} {:>37}\n".format(*header),
+            u"-" * 80
         ]
 
         for n, contraction in enumerate(self.contraction_list):
             inds, idx_rm, einsum_str, remaining, do_blas = contraction
             remaining_str = ",".join(remaining) + "->" + self.output_subscript
             path_run = (self.scale_list[n], do_blas, einsum_str, remaining_str)
-            path_print.append("\n{:>4} {:>14} {:>22} {:>37}".format(*path_run))
+            path_print.append(u"\n{:>4} {:>14} {:>22} {:>37}".format(*path_run))
 
         return "".join(path_print)
 
@@ -99,6 +99,8 @@ def contract_path(*operands, **kwargs):
         - 'optimal' An algorithm that tries all possible ways of
           contracting the listed tensors. Scales exponentially with
           the number of terms in the contraction.
+        - 'eager' An approximate algorithm that is cheaper than 'greedy' for
+          very large contractions.
         - 'auto' Use the optimal approach for a small number of terms (currently
           4 or less), else use the greedy approach.
 
@@ -253,6 +255,8 @@ def contract_path(*operands, **kwargs):
         path = paths.optimal(input_sets, output_set, dimension_dict, memory_arg)
     elif path_type in ("greedy", "opportunistic", "auto"):
         path = paths.greedy(input_sets, output_set, dimension_dict, memory_arg)
+    elif path_type == 'eager':
+        path = paths.eager(input_sets, output_set, dimension_dict)
     else:
         raise KeyError("Path name '{}' not found".format(path_type))
 
@@ -390,6 +394,8 @@ def contract(*operands, **kwargs):
         - 'optimal' An algorithm that tries all possible ways of
           contracting the listed tensors. Scales exponentially with
           the number of terms in the contraction.
+        - 'eager' An approximate algorithm that is cheaper than 'greedy'
+          for very large contractions.
 
     memory_limit : {None, int, 'max_input'} (default: None)
         Give the upper bound of the largest intermediate tensor contract will build.
diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py
index 56a96ea..8ed36d0 100644
--- a/opt_einsum/paths.py
+++ b/opt_einsum/paths.py
@@ -1,11 +1,18 @@
 """
 Contains the path technology behind opt_einsum in addition to several path helpers
 """
+from __future__ import absolute_import, division, print_function
+
+import heapq
 import itertools
+from collections import defaultdict
+
+import numpy as np
 
 from . import helpers
 
-__all__ = ["optimal", "greedy"] 
+__all__ = ["optimal", "greedy", "eager"]
+
 
 def optimal(input_sets, output_set, idx_dict, memory_limit):
     """
@@ -272,3 +279,210 @@ def greedy(input_sets, output_set, idx_dict, memory_limit):
         path_cost += best[0][1]
 
     return path
+
+
+def ssa_to_linear(ssa_path):
+    """
+    Convert a path with static single assignment ids to a path with recycled
+    linear ids. For example::
+
+        >>> ssa_to_linear([(0, 3), (2, 4), (1, 5)])
+        [(0, 3), (1, 2), (0, 1)]
+    """
+    ids = np.arange(1 + max(map(max, ssa_path)), dtype=np.int32)
+    path = []
+    for ssa_ids in ssa_path:
+        path.append(tuple(int(ids[ssa_id]) for ssa_id in ssa_ids))
+        for ssa_id in ssa_ids:
+            ids[ssa_id:] -= 1
+    return path
+
+
+def linear_to_ssa(path):
+    """
+    Convert a path with recycled linear ids to a path with static single
+    assignment ids. For example::
+
+        >>> linear_to_ssa([(0, 3), (1, 2), (0, 1)])
+        [(0, 3), (2, 4), (1, 5)]
+    """
+    num_inputs = sum(map(len, path)) - len(path) + 1
+    linear_to_ssa = list(range(num_inputs))
+    new_ids = itertools.count(num_inputs)
+    ssa_path = []
+    for ids in path:
+        ssa_path.append(tuple(linear_to_ssa[id_] for id_ in ids))
+        for id_ in sorted(ids, reverse=True):
+            del linear_to_ssa[id_]
+        linear_to_ssa.append(next(new_ids))
+    return ssa_path
+
+
+def _get_candidate(output, sizes, remaining, footprints, dim_ref_counts, k1, k2):
+    either = k1 | k2
+    two = k1 & k2
+    one = either - two
+    k12 = (either & output) | (two & dim_ref_counts[3]) | (one & dim_ref_counts[2])
+    cost = helpers.compute_size_by_dict(k12, sizes) - footprints[k1] - footprints[k2]
+    id1 = remaining[k1]
+    id2 = remaining[k2]
+    if id1 > id2:
+        k1, id1, k2, id2 = k2, id2, k1, id1
+    cost = cost, id2, id1  # break ties to ensure determinism
+    return cost, k1, k2, k12
+
+
+def _push_candidate(output, sizes, remaining, footprints, dim_ref_counts, k1, k2s, queue):
+    candidate = min(_get_candidate(output, sizes, remaining, footprints, dim_ref_counts, k1, k2)
+                    for k2 in k2s)
+    heapq.heappush(queue, candidate)
+
+
+def _update_ref_counts(dim_to_keys, dim_ref_counts, dims):
+    for dim in dims:
+        count = len(dim_to_keys[dim])
+        if count <= 1:
+            dim_ref_counts[2].discard(dim)
+            dim_ref_counts[3].discard(dim)
+        elif count == 2:
+            dim_ref_counts[2].add(dim)
+            dim_ref_counts[3].discard(dim)
+        else:
+            dim_ref_counts[2].add(dim)
+            dim_ref_counts[3].add(dim)
+
+
+def _ssa_optimize(inputs, output, sizes):
+    """
+    This has an interface similar to :func:`eager` but produces a path with
+    static single assignment ids rather than recycled linear ids.
+    SSA ids are cheaper to work with and easier to reason about.
+    """
+    if len(inputs) == 1:
+        # Perform a single contraction to match output shape.
+        return [(0,)]
+
+    # A dim that is common to all tensors might as well be an output dim, since it
+    # cannot be contracted until the final step. This avoids an expensive all-pairs
+    # comparison to search for possible contractions at each step, leading to speedup
+    # in many practical problems where all tensors share a common batch dimension.
+    inputs = list(map(frozenset, inputs))
+    output = frozenset(output) | frozenset.intersection(*inputs)
+
+    # Deduplicate shapes by eagerly computing Hadamard products.
+    remaining = {}  # key -> ssa_id
+    ssa_ids = itertools.count(len(inputs))
+    ssa_path = []
+    for ssa_id, key in enumerate(inputs):
+        if key in remaining:
+            ssa_path.append((remaining[key], ssa_id))
+            remaining[key] = next(ssa_ids)
+        else:
+            remaining[key] = ssa_id
+
+    # Keep track of possible contraction dims.
+    dim_to_keys = defaultdict(set)
+    for key in remaining:
+        for dim in key - output:
+            dim_to_keys[dim].add(key)
+
+    # Keep track of the number of tensors using each dim; when the dim is no longer
+    # used it can be contracted. Since we specialize to binary ops, we only care about
+    # ref counts of >=2 or >=3.
+    dim_ref_counts = {
+        count: set(dim for dim, keys in dim_to_keys.items() if len(keys) >= count) - output
+        for count in [2, 3]}
+
+    # Compute separable part of the objective function for contractions.
+    footprints = {key: helpers.compute_size_by_dict(key, sizes) for key in remaining}
+
+    # Find initial candidate contractions.
+    queue = []
+    for dim, keys in dim_to_keys.items():
+        keys = sorted(keys, key=remaining.__getitem__)
+        for i, k1 in enumerate(keys[:-1]):
+            k2s = keys[1 + i:]
+            _push_candidate(output, sizes, remaining, footprints, dim_ref_counts, k1, k2s, queue)
+
+    # Greedily contract pairs of tensors.
+    while queue:
+        cost, k1, k2, k12 = heapq.heappop(queue)
+        if k1 not in remaining or k2 not in remaining:
+            continue  # candidate is obsolete
+
+        ssa_id1 = remaining.pop(k1)
+        ssa_id2 = remaining.pop(k2)
+        for dim in k1 - output:
+            dim_to_keys[dim].remove(k1)
+        for dim in k2 - output:
+            dim_to_keys[dim].remove(k2)
+        ssa_path.append((ssa_id1, ssa_id2))
+        if k12 in remaining:
+            ssa_path.append((remaining[k12], next(ssa_ids)))
+        else:
+            for dim in k12 - output:
+                dim_to_keys[dim].add(k12)
+        remaining[k12] = next(ssa_ids)
+        _update_ref_counts(dim_to_keys, dim_ref_counts, k1 | k2 - output)
+        footprints[k12] = helpers.compute_size_by_dict(k12, sizes)
+
+        # Find new candidate contractions.
+        k1 = k12
+        k2s = set(k2 for dim in k1 for k2 in dim_to_keys[dim])
+        k2s.discard(k1)
+        if k2s:
+            _push_candidate(output, sizes, remaining, footprints, dim_ref_counts, k1, k2s, queue)
+
+    # Greedily compute pairwise outer products.
+    queue = [(helpers.compute_size_by_dict(key & output, sizes), ssa_id, key)
+             for key, ssa_id in remaining.items()]
+    heapq.heapify(queue)
+    _, ssa_id1, k1 = heapq.heappop(queue)
+    while queue:
+        _, ssa_id2, k2 = heapq.heappop(queue)
+        ssa_path.append((min(ssa_id1, ssa_id2), max(ssa_id1, ssa_id2)))
+        k12 = (k1 | k2) & output
+        cost = helpers.compute_size_by_dict(k12, sizes)
+        ssa_id12 = next(ssa_ids)
+        _, ssa_id1, k1 = heapq.heappushpop(queue, (cost, ssa_id12, k12))
+
+    return ssa_path
+
+
+def eager(inputs, output, idx_dict):
+    """
+    Finds the path by a three stage algorithm:
+
+    1. Eagerly compute Hadamard products.
+    2. Greedily compute contractions to maximize ``removed_size``
+    3. Greedily compute outer products.
+
+    This algorithm scales quadratically with respect to the
+    maximum number of elements sharing a common dim.
+
+    Parameters
+    ----------
+    input_sets : list
+        List of sets that represent the lhs side of the einsum subscript
+    output_set : set
+        Set that represents the rhs side of the overall einsum subscript
+    idx_dict : dictionary
+        Dictionary of index sizes
+    memory_limit : int
+        The maximum number of elements in a temporary array
+
+    Returns
+    -------
+    path : list
+        The contraction order (a list of tuples of ints).
+
+    Examples
+    --------
+    >>> isets = [set('abd'), set('ac'), set('bdc')]
+    >>> oset = set('')
+    >>> idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
+    >>> eager(isets, oset, idx_sizes)
+    [(0, 2), (0, 1)]
+    """
+    ssa_path = _ssa_optimize(inputs, output, idx_dict)
+    return ssa_to_linear(ssa_path)


Test patch:
diff --git a/opt_einsum/tests/test_contract.py b/opt_einsum/tests/test_contract.py
index 22900c8..dbbc39e 100644
--- a/opt_einsum/tests/test_contract.py
+++ b/opt_einsum/tests/test_contract.py
@@ -6,6 +6,7 @@ import numpy as np
 import pytest
 
 from opt_einsum import compat, contract, contract_path, helpers, contract_expression
+from opt_einsum.paths import linear_to_ssa, ssa_to_linear
 
 tests = [
     # Test hadamard-like products
@@ -97,14 +98,12 @@ tests = [
 
 
 @pytest.mark.parametrize("string", tests)
-def test_compare(string):
+@pytest.mark.parametrize("optimize", ['greedy', 'optimal', 'eager'])
+def test_compare(optimize, string):
     views = helpers.build_views(string)
 
     ein = contract(string, *views, optimize=False, use_blas=False)
-    opt = contract(string, *views, optimize='greedy', use_blas=False)
-    assert np.allclose(ein, opt)
-
-    opt = contract(string, *views, optimize='optimal', use_blas=False)
+    opt = contract(string, *views, optimize=optimize, use_blas=False)
     assert np.allclose(ein, opt)
 
 
@@ -116,7 +115,8 @@ def test_drop_in_replacement(string):
 
 
 @pytest.mark.parametrize("string", tests)
-def test_compare_greek(string):
+@pytest.mark.parametrize("optimize", ['greedy', 'optimal', 'eager'])
+def test_compare_greek(optimize, string):
     views = helpers.build_views(string)
 
     ein = contract(string, *views, optimize=False, use_blas=False)
@@ -124,27 +124,23 @@ def test_compare_greek(string):
     # convert to greek
     string = ''.join(compat.get_char(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)
-
-    opt = contract(string, *views, optimize='optimal', use_blas=False)
+    opt = contract(string, *views, optimize=optimize, use_blas=False)
     assert np.allclose(ein, opt)
 
 
 @pytest.mark.parametrize("string", tests)
-def test_compare_blas(string):
+@pytest.mark.parametrize("optimize", ['greedy', 'optimal', 'eager'])
+def test_compare_blas(optimize, string):
     views = helpers.build_views(string)
 
     ein = contract(string, *views, optimize=False)
-    opt = contract(string, *views, optimize='greedy')
-    assert np.allclose(ein, opt)
-
-    opt = contract(string, *views, optimize='optimal')
+    opt = contract(string, *views, optimize=optimize)
     assert np.allclose(ein, opt)
 
 
 @pytest.mark.parametrize("string", tests)
-def test_compare_blas_greek(string):
+@pytest.mark.parametrize("optimize", ['greedy', 'optimal', 'eager'])
+def test_compare_blas_greek(optimize, string):
     views = helpers.build_views(string)
 
     ein = contract(string, *views, optimize=False)
@@ -152,10 +148,7 @@ def test_compare_blas_greek(string):
     # convert to greek
     string = ''.join(compat.get_char(ord(c) + 848) if c not in ',->.' else c for c in string)
 
-    opt = contract(string, *views, optimize='greedy')
-    assert np.allclose(ein, opt)
-
-    opt = contract(string, *views, optimize='optimal')
+    opt = contract(string, *views, optimize=optimize)
     assert np.allclose(ein, opt)
 
 
@@ -177,7 +170,7 @@ def test_printing():
 
 
 @pytest.mark.parametrize("string", tests)
-@pytest.mark.parametrize("optimize", ['greedy', 'optimal'])
+@pytest.mark.parametrize("optimize", ['greedy', 'optimal', 'eager'])
 @pytest.mark.parametrize("use_blas", [False, True])
 @pytest.mark.parametrize("out_spec", [False, True])
 def test_contract_expressions(string, optimize, use_blas, out_spec):
@@ -243,3 +236,12 @@ def test_rand_equation(optimize, n, reg, n_out):
     actual = contract(eq, *views, optimize=optimize)
 
     assert np.allclose(expected, actual)
+
+
+@pytest.mark.parametrize('equation', tests)
+def test_linear_vs_ssa(equation):
+    views = helpers.build_views(equation)
+    linear_path, _ = contract_path(equation, *views)
+    ssa_path = linear_to_ssa(linear_path)
+    linear_path2 = ssa_to_linear(ssa_path)
+    assert linear_path2 == linear_path
diff --git a/opt_einsum/tests/test_paths.py b/opt_einsum/tests/test_paths.py
index dc61645..705f829 100644
--- a/opt_einsum/tests/test_paths.py
+++ b/opt_einsum/tests/test_paths.py
@@ -26,12 +26,16 @@ explicit_path_tests = {
 }
 
 path_edge_tests = [
+    ['eager', 'eb,cb,fb->cef', ((0, 2), (0, 1))],
     ['greedy', 'eb,cb,fb->cef', ((0, 2), (0, 1))],
     ['optimal', 'eb,cb,fb->cef', ((0, 2), (0, 1))],
+    ['eager', 'dd,fb,be,cdb->cef', ((0, 3), (0, 1), (0, 1))],
     ['greedy', 'dd,fb,be,cdb->cef', ((0, 3), (0, 1), (0, 1))],
     ['optimal', 'dd,fb,be,cdb->cef', ((0, 3), (0, 1), (0, 1))],
+    ['eager', 'bca,cdb,dbf,afc->', ((1, 2), (0, 2), (0, 1))],
     ['greedy', 'bca,cdb,dbf,afc->', ((1, 2), (0, 2), (0, 1))],
     ['optimal', 'bca,cdb,dbf,afc->', ((1, 2), (0, 2), (0, 1))],
+    ['eager', 'dcc,fce,ea,dbf->ab', ((1, 2), (0, 1), (0, 1))],
     ['greedy', 'dcc,fce,ea,dbf->ab', ((1, 2), (0, 1), (0, 1))],
     ['optimal', 'dcc,fce,ea,dbf->ab', ((1, 2), (0, 2), (0, 1))],
 ]
@@ -172,10 +176,11 @@ def test_greedy_edge_cases():
     assert check_path(path, [(0, 1), (0, 2), (0, 1)])
 
 
-def test_can_optimize_outer_products():
+@pytest.mark.parametrize("optimize", ['greedy', 'eager'])
+def test_can_optimize_outer_products(optimize):
     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, optimize='greedy')[0] == [(2, 3), (0, 2), (0, 1)]
+    assert oe.contract_path("ab,cd,ef,fg", a, b, c, d, optimize=optimize)[0] == [(2, 3), (0, 2), (0, 1)]
 
 
 @pytest.mark.parametrize('num_symbols', [2, 3, 26, 26 + 26, 256 - 140, 300])