Repository: dgasmith/opt_einsum
Pull request #54 (2018-08-25 15:45:08)

v2.2 doc cleanup
## Description
General docs and consistency checks enumerated as follows:
- [x] Runs YAPF over the code.
- [x] Appends new version to the changelog.
- [x] Expands comments on `greedy` paths.
- [x] Adds small mention about shared intermediate thread sharing.
- [x] Adds docstrings and examples where needed.
- [x] Uses `__all__` consistently throughout the code to ensure globals are correctly protected.

## Status
- [x] Ready to go


Patch:
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 34e34c2..fbfdf21 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -6,11 +6,37 @@ Changelog
 
 New Features
 ++++++++++++
+- (:pr:`48`) Intermediates can now be shared between contractions, see here for more details.
+- (:pr:`53`) Intermediate caching is thread safe.
 
-- Improved documentation.
+Enhancements
+++++++++++++
+- (:pr:`48`) Expressions are now mapped to non-unicode index set so that unicode input is support for all backends.
+- (:pr:`54`) General documenation update.
+
+Bug fixes
++++++++++
+- (:pr:`41`) PyTorch indices are mapped back to a small a-z subset valid for PyTorch's einsum implementation.
+
+2.1.3 / 2018-8-23
+-----------------
+
+Bug fixes
++++++++++
+
+- Fixes unicode issue for large numbers of tensors in Python 2.7.
+- Fixes unicode install bug in README.md.
+
+2.1.2 / 2018-8-16
+-----------------
+
+Bug fixes
++++++++++
+
+- Ensures `versioneer.py` is in MANIFEST.in for a clean pip install.
 
 
-2.1.1 / 2018-8-16
+2.1.1 / 2018-8-15
 -----------------
 
 Bug fixes
@@ -71,7 +97,7 @@ New Features
 Bug Fixes
 +++++++++
 
-- A error effecting cases where opt_einsum mistook broadcasting operations for matrix multiply has been fixed.
+- An error affecting cases where opt_einsum mistook broadcasting operations for matrix multiply has been fixed.
 - Most error messages are now more expressive.
 
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 49c3b2c..f5b6eb9 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -43,6 +43,7 @@ extensions = [
     'sphinx.ext.viewcode',
     'sphinx.ext.napoleon',
     'sphinx.ext.intersphinx',
+    'sphinx.ext.extlinks',
 ]
 
 napoleon_google_docstring = False
diff --git a/docs/source/ex_dask_reusing_intermediaries.rst b/docs/source/ex_dask_reusing_intermediaries.rst
index e573314..8c9b4f0 100644
--- a/docs/source/ex_dask_reusing_intermediaries.rst
+++ b/docs/source/ex_dask_reusing_intermediaries.rst
@@ -4,10 +4,10 @@ Reusing Intermediaries with Dask
 `Dask <https://dask.pydata.org/>`_ provides a computational framework where
 arrays and the computations on them are built up into a 'task graph' before
 computation. Since :mod:`opt_einsum` is compatible with ``dask`` arrays this
-means that mutiple contractions can be built into the same task graph, which
+means that multiple contractions can be built into the same task graph, which
 then automatically reuses any shared arrays and contractions.
 
-For example imagine the two expressions:
+For example, imagine the two expressions:
 
 .. code-block:: python
 
@@ -31,7 +31,7 @@ First, let's set up some ``numpy`` arrays:
     >>> np_ops1 = [np_arrays[s] for s in terms1]
     >>> np_ops2 = [np_arrays[s] for s in terms2]
 
-Normally we would just compute these separately:
+Typically we would compute these expressions separately:
 
 .. code-block:: python
 
@@ -41,7 +41,7 @@ Normally we would just compute these separately:
     >>> oe.contract(contraction2, *np_ops2)
     array(-75.55902751)
 
-But if we use dask arrays we can combine the two operations, so let's set those
+However, if we use dask arrays we can combine the two operations, so let's set those
 up:
 
 .. code-block:: python
@@ -72,17 +72,17 @@ Note ``chunks`` is a required argument relating to how the arrays are stored (se
     >>> dy
     Delayed('list-3af82335-b75e-47d6-b800-68490fc865fd')
 
-As suggested by the name ``Delayed``, we just have a placeholder for the result
-so far. When we want to actually *perform* the computation we can call:
+As suggested by the name ``Delayed``, we have a placeholder for the result
+so far. When we want to *perform* the computation we can call:
 
 .. code-block:: python
 
     >>> dy.compute()
     [114.78314052155015, -75.55902750513113]
 
-Which matches the numpy result. The computation can even be handled by various
+The above matches the canonical numpy result. The computation can even be handled by various
 schedulers - see `scheduling <http://dask.pydata.org/en/latest/scheduling.html>`_.
-Finally, to check we really are reusing intermediaries we can view the task
+Finally, to check we are reusing intermediaries, we can view the task
 graph generated for the computation:
 
 .. code-block:: python
@@ -90,3 +90,6 @@ graph generated for the computation:
     >>> dy.visualize(optimize_graph=True)
 
 .. image:: ex_dask_reuse_graph.png
+
+.. note::
+    For sharing intermediates with other backends see `Sharing Intermediates <sharing_intermediates.html>`_. Dask graphs are particularly useful for reusing intermediates beyond just contractions and can allow additional parallelization.
diff --git a/docs/source/greedy_path.rst b/docs/source/greedy_path.rst
index 8ff1621..4a32ab7 100644
--- a/docs/source/greedy_path.rst
+++ b/docs/source/greedy_path.rst
@@ -9,5 +9,51 @@ This algorithm has proven to be quite successful for general production and only
 Fortunately, these often only lead to increases in prefactor than missing the optimal scaling.
 
 The ``greedy`` approach scales like N^2 rather than factorially, making ``greedy`` much more suitable for large numbers of contractions where the lower prefactor helps decrease latency.
-As :mod:`opt_einsum` can handle an arbitrary number of indices the low scaling is especially important for very large contraction networks.
+As :mod:`opt_einsum` can handle an arbitrary number of indices the low scaling is necessary for extensive contraction networks.
 The ``greedy`` functionality is provided by :func:`~opt_einsum.paths.greedy`.
+
+Optimal Scaling Misses
+----------------------
+The greedy algorithm, while inexpensive, can occasionally miss optimal scaling in some circumstances as seen below. The ``greedy`` algorithm prioritizes expressions which remove the largest indices first, in this particular case this is the incorrect choice and it is difficult for any heuristic algorithm to "see ahead" as would be needed here.
+
+It should be stressed these cases are quite rare and by default ``contract`` uses the ``optimal`` path for four and fewer inputs as the cost of evaluating the ``optimal`` path is similar to that of the ``greedy`` path.
+
+.. code:: python
+
+    >>> M = np.random.rand(35, 37, 59)
+    >>> A = np.random.rand(35, 51, 59)
+    >>> B = np.random.rand(37, 51, 51, 59)
+    >>> C = np.random.rand(59, 27)
+
+    >>> path, desc = oe.contract_path('xyf,xtf,ytpf,fr->tpr', M, A, B, C, path="greedy")
+    >>> print(desc)
+      Complete contraction:  xyf,xtf,ytpf,fr->tpr
+             Naive scaling:  6
+         Optimized scaling:  5
+          Naive FLOP count:  2.146e+10
+      Optimized FLOP count:  4.165e+08
+       Theoretical speedup:  51.533
+      Largest intermediate:  5.371e+06 elements
+    --------------------------------------------------------------------------------
+    scaling        BLAS                current                             remaining
+    --------------------------------------------------------------------------------
+       5          False         ytpf,xyf->tpfx                      xtf,fr,tpfx->tpr
+       4          False          tpfx,xtf->tpf                           fr,tpf->tpr
+       4           GEMM            tpf,fr->tpr                              tpr->tpr
+
+    >>> path, desc = oe.contract_path('xyf,xtf,ytpf,fr->tpr', M, A, B, C, path="optimal")
+    >>> print(desc)
+
+      Complete contraction:  xyf,xtf,ytpf,fr->tpr
+             Naive scaling:  6
+         Optimized scaling:  4
+          Naive FLOP count:  2.146e+10
+      Optimized FLOP count:  2.744e+07
+       Theoretical speedup:  782.283
+      Largest intermediate:  1.535e+05 elements
+    --------------------------------------------------------------------------------
+    scaling        BLAS                current                             remaining
+    --------------------------------------------------------------------------------
+       4          False           xtf,xyf->tfy                      ytpf,fr,tfy->tpr
+       4          False          tfy,ytpf->tfp                           fr,tfp->tpr
+       4           TDOT            tfp,fr->tpr                              tpr->tpr
\ No newline at end of file
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 869a28f..82776a1 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -2,7 +2,7 @@
 opt_einsum
 ==========
 
-:func:`~numpy.einsum` is a very powerful function for contracting tensors of arbitrary
+:func:`~numpy.einsum` is a powerful function for contracting tensors of arbitrary
 dimension and index. However, it is only optimized to contract two terms
 at a time resulting in non-optimal scaling.
 
@@ -32,8 +32,8 @@ Consider two different algorithms:
         return K
 
 The einsum function does not consider building intermediate arrays;
-therefore, helping einsum out by building these intermediate arrays can result
-in a considerable cost savings even for small N (N=10):
+therefore, helping einsum out by creating these intermediate arrays can result
+in considerable cost savings even for small N (N=10):
 
 .. code:: python
 
@@ -46,14 +46,14 @@ in a considerable cost savings even for small N (N=10):
     %timeit optimized(I, C)
     1000 loops, best of 3: 445 µs per loop
 
-The index transformation is a well known contraction that leads to
+The index transformation is a well-known contraction that leads to
 straightforward intermediates. This contraction can be further
 complicated by considering that the shape of the C matrices need not be
-the same, in this case the ordering in which the indices are transformed
-matters greatly. Logic can be built that optimizes the ordering;
+the same, in this case, the ordering in which the indices are transformed
+matters significantly. Logic can be built that optimizes the order;
 however, this is a lot of time and effort for a single expression.
 
-The opt_einsum package is a drop in replacement for the ``np.einsum`` function
+The opt_einsum package is a drop-in replacement for the ``np.einsum`` function
 and can handle all of the logic for you:
 
 .. code:: python
@@ -70,8 +70,7 @@ and can handle all of the logic for you:
     %timeit contract('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
     100 loops, best of 3: 16.2 ms per loop
 
-The above will automatically find the optimal contraction order, in this case
-identical to that of the optimized function above, and compute the products for
+The above will automatically find the optimal contraction order, in this case, identical to that of the optimized function above, and compute the products for
 you. In this case, it even uses `np.dot` under the hood to exploit any vendor
 BLAS functionality that your NumPy build has!
 
diff --git a/docs/source/optimal_path.rst b/docs/source/optimal_path.rst
index 3db5546..76f5400 100644
--- a/docs/source/optimal_path.rst
+++ b/docs/source/optimal_path.rst
@@ -5,11 +5,11 @@ The Optimal Path
 The most optimal path can be found by searching through every possible way to contract the tensors together, this includes all combinations with the new intermediate tensors as well.
 While this algorithm scales like N! and can often become more costly to compute than the unoptimized contraction itself, it provides an excellent benchmark.
 The function that computes this path in opt_einsum is called :func:`~opt_einsum.paths.optimal` and works by iteratively finding every possible combination of pairs to contract in the current list of tensors.
-This is iterated until all tensors are contracted together. The resulting paths are then sorted by total flop cost and the lowest one is chosen.
+This is iterated until all tensors are contracted together. The resulting paths are then sorted by total FLOP cost and the lowest one is chosen.
 This algorithm runs in about 1 second for 7 terms, 15 seconds for 8 terms, and 480 seconds for 9 terms limiting its overall usefulness for a large number of terms.
-By considering limited memory this can be sieved and can reduce the cost of computing the optimal function by an order of magnitude or more.
+By considering limited memory, this can be sieved and can reduce the cost of computing the optimal function by an order of magnitude or more.
 
-Lets look at an example:
+Let us look at an example:
 
 .. code:: python
 
@@ -43,4 +43,4 @@ As all contractions only have two remaining input sets the only possible contrac
       (3168,  [(0, 2), (0, 1)], [set(['b', 'd'])  ]),
       (19872, [(1, 2), (0, 1)], [set(['b', 'd'])  ])]
 
-The final contraction cost is computed and we choose the second path from the list as the overall cost is the lowest.
+The final contraction cost is computed, and we choose the second path from the list as the overall cost is the lowest.
diff --git a/docs/source/reusing_paths.rst b/docs/source/reusing_paths.rst
index d630582..4661fbe 100644
--- a/docs/source/reusing_paths.rst
+++ b/docs/source/reusing_paths.rst
@@ -2,7 +2,7 @@
 Reusing Paths
 =============
 
-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 :func:`~opt_einsum.contract_expression` with the contraction string and the shapes of the tensors generates a :class:`~opt_einsum.contract.ContractExpression` which can then be repeatedly called with any matching set of arrays. For example:
+If you expect to use a particular contraction repeatedly, it can make things simpler and more efficient not to compute the path each time. Instead, supplying :func:`~opt_einsum.contract_expression` with the contraction string and the shapes of the tensors generates a :class:`~opt_einsum.contract.ContractExpression` which can then be repeatedly called with any matching set of arrays. For example:
 
 .. code:: python
 
@@ -67,7 +67,7 @@ supply the actual arrays rather than just the shapes to
     <ContractExpression('ij,[jk,kl,lm],mn->ni', constants=[1, 2, 3])>
 
 The expression now only takes the remaining two arrays as arguments (the
-tensors with ``'ij'`` and ``'mn'`` indices), and will store as many resuable
+tensors with ``'ij'`` and ``'mn'`` indices), and will store as many reusable
 constant contractions as possible.
 
 .. code:: python
@@ -95,9 +95,8 @@ two contractions to compute the output.
 
 .. note::
 
-    The constant part of an expression is lazily generated upon first call
-    (specific to each backend), though it can also be explicitly built by
-    calling :meth:`~opt_einsum.contract.ContractExpression.evaluate_constants`.
+    The constant part of an expression is lazily generated upon the first call
+    (specific to each backend), though it can also be explicitly built by calling:meth:`~opt_einsum.contract.ContractExpression.evaluate_constants`.
 
 We can confirm the advantage of using expressions and constants by timing the
 following scenarios, first setting
@@ -132,4 +131,4 @@ contractions.
 We also note that even if there are *no* constant contractions to perform, it
 can be very advantageous to specify constant tensors for particular backends.
 For instance, if a GPU backend is used, the constant tensors will be kept on
-the device rather than being transfered each time.
+the device rather than being transferred each time.
diff --git a/docs/source/sharing_intermediates.rst b/docs/source/sharing_intermediates.rst
index 345a2c4..da1d619 100644
--- a/docs/source/sharing_intermediates.rst
+++ b/docs/source/sharing_intermediates.rst
@@ -16,7 +16,7 @@ For example, suppose we want to compute marginals at each point in a factor chai
                  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
+To share this computation, we can perform all contractions in a shared context:
 
 .. code:: python
 
@@ -39,3 +39,6 @@ If it is difficult to fit your code into a context, you can instead save the sha
     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.
+
+.. note::
+    By default a cache is thread safe, to share intermediates between threads explicitly pass the same cache to each thread.
diff --git a/opt_einsum/__init__.py b/opt_einsum/__init__.py
index 1db2b64..324bcdc 100644
--- a/opt_einsum/__init__.py
+++ b/opt_einsum/__init__.py
@@ -2,15 +2,15 @@
 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
-
+from . import paths
 # Handle versioneer
 from ._version import get_versions
+from .contract import contract, contract_path, contract_expression
+from .parser import get_symbol
+from .sharing import shared_intermediates
+
 versions = get_versions()
 __version__ = versions['version']
 __git_revision__ = versions['full-revisionid']
diff --git a/opt_einsum/backends/__init__.py b/opt_einsum/backends/__init__.py
index 33e5362..1e8915b 100644
--- a/opt_einsum/backends/__init__.py
+++ b/opt_einsum/backends/__init__.py
@@ -1,7 +1,13 @@
-from .dispatch import (get_func, has_einsum, build_expression, evaluate_constants, CONVERT_BACKENDS)
+"""
+Compute backends for opt_einsum.
+"""
+
+from .dispatch import (get_func, has_einsum, build_expression, evaluate_constants, has_backend)
+
+# Backends
+from .cupy import to_cupy
 from .tensorflow import to_tensorflow
 from .theano import to_theano
-from .cupy import to_cupy
 from .torch import to_torch
 
 __all__ = [
diff --git a/opt_einsum/backends/cupy.py b/opt_einsum/backends/cupy.py
index ce7bf6b..21f3d30 100644
--- a/opt_einsum/backends/cupy.py
+++ b/opt_einsum/backends/cupy.py
@@ -3,9 +3,13 @@ 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
 
+__all__ = ["to_cupy", "build_expression", "evaluate_constants"]
+
 
 @to_backend_cache_wrap
 def to_cupy(array):  # pragma: no cover
@@ -31,4 +35,4 @@ def evaluate_constants(const_arrays, expr):  # pragma: no cover
     """Convert constant arguments to cupy arrays, and perform any possible
     constant contractions.
     """
-    return expr(*[to_cupy(x) for x in const_arrays], backend='cupy', evaluate_constants=True)
+    return expr(* [to_cupy(x) for x in const_arrays], backend='cupy', evaluate_constants=True)
diff --git a/opt_einsum/backends/dispatch.py b/opt_einsum/backends/dispatch.py
index f766f3f..7d84fd4 100644
--- a/opt_einsum/backends/dispatch.py
+++ b/opt_einsum/backends/dispatch.py
@@ -9,9 +9,11 @@ import importlib
 import numpy
 
 from . import cupy as _cupy
-from . import torch as _torch
-from . import theano as _theano
 from . import tensorflow as _tensorflow
+from . import theano as _theano
+from . import torch as _torch
+
+__all__ = ["get_func", "has_einsum", "build_expression", "evaluate_constants", "has_backend"]
 
 # known non top-level imports
 _aliases = {
@@ -80,7 +82,6 @@ CONVERT_BACKENDS = {
     'torch': _torch.build_expression,
 }
 
-
 EVAL_CONSTS_BACKENDS = {
     'tensorflow': _tensorflow.evaluate_constants,
     'theano': _theano.evaluate_constants,
@@ -101,3 +102,9 @@ def evaluate_constants(backend, arrays, expr):
     the contraction of ``expr`` with these as possible.
     """
     return EVAL_CONSTS_BACKENDS[backend](arrays, expr)
+
+
+def has_backend(backend):
+    """Checks if the backend is known.
+    """
+    return backend.lower() in CONVERT_BACKENDS
diff --git a/opt_einsum/backends/tensorflow.py b/opt_einsum/backends/tensorflow.py
index a6d218b..3ac2967 100644
--- a/opt_einsum/backends/tensorflow.py
+++ b/opt_einsum/backends/tensorflow.py
@@ -3,8 +3,10 @@ Required functions for optimized contractions of numpy arrays using tensorflow.
 """
 
 from __future__ import absolute_import
+
 import numpy as np
 
+__all__ = ["to_tensorflow", "build_expression", "evaluate_constants"]
 
 _CACHED_TF_DEVICE = None
 
@@ -52,6 +54,7 @@ def to_tensorflow(array, constant=False):
 
 # Standard graph mode
 
+
 def build_expression_graph(arrays, expr):
     """Build a tensorflow function based on ``arrays`` and ``expr``.
     """
@@ -89,9 +92,11 @@ def evaluate_constants_graph(const_arrays, expr):
 
 # Eager execution mode
 
+
 def build_expression_eager(_, expr):
     """Build a eager tensorflow function based on ``arrays`` and ``expr``.
     """
+
     def tensorflow_eager_contract(*arrays):
         return expr._contract([to_tensorflow(x) for x in arrays], backend='tensorflow').numpy()
 
@@ -102,11 +107,12 @@ def evaluate_constants_eager(const_arrays, expr):
     """Convert constant arguments to tensorflow_eager arrays, and perform any
     possible constant contractions.
     """
-    return expr(*[to_tensorflow(x) for x in const_arrays], backend='tensorflow', evaluate_constants=True)
+    return expr(* [to_tensorflow(x) for x in const_arrays], backend='tensorflow', evaluate_constants=True)
 
 
 # Dispatch to eager or graph mode
 
+
 def build_expression(arrays, expr):
     _, _, eager = _get_tensorflow_and_device()
     fn = build_expression_eager if eager else build_expression_graph
diff --git a/opt_einsum/backends/theano.py b/opt_einsum/backends/theano.py
index 80c0bef..65817f3 100644
--- a/opt_einsum/backends/theano.py
+++ b/opt_einsum/backends/theano.py
@@ -3,8 +3,11 @@ Required functions for optimized contractions of numpy arrays using theano.
 """
 
 from __future__ import absolute_import
+
 import numpy as np
 
+__all__ = ["to_theano", "build_expression", "evaluate_constants"]
+
 
 def to_theano(array, constant=False):
     """Convert a numpy array to ``theano.tensor.TensorType`` instance.
@@ -33,7 +36,7 @@ def build_expression(arrays, expr):
     graph = theano.function(graph_ins, out_var)
 
     def theano_contract(*arrays):
-        return graph(*[x for x in arrays if not isinstance(x, theano.tensor.TensorConstant)])
+        return graph(* [x for x in arrays if not isinstance(x, theano.tensor.TensorConstant)])
 
     return theano_contract
 
diff --git a/opt_einsum/backends/torch.py b/opt_einsum/backends/torch.py
index 124d87b..90edec1 100644
--- a/opt_einsum/backends/torch.py
+++ b/opt_einsum/backends/torch.py
@@ -3,13 +3,18 @@ Required functions for optimized contractions of numpy arrays using pytorch.
 """
 
 from __future__ import absolute_import
+
 import numpy as np
 
-from ..parser import convert_to_valid_einsum_chars, einsum_symbols_base
+from ..parser import convert_to_valid_einsum_chars
 from ..sharing import to_backend_cache_wrap
 
+__all__ = ["transpose", "einsum", "tensordot", "to_torch", "build_expression", "evaluate_constants"]
+
 _TORCH_DEVICE = None
 
+_torch_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
+
 
 def _get_torch_and_device():
     global _TORCH_DEVICE
@@ -62,7 +67,7 @@ def tensordot(x, y, axes=2):
     out_ix = []
 
     # fill in repeated indices
-    available_ix = iter(einsum_symbols_base)
+    available_ix = iter(_torch_symbols_base)
     for ax1, ax2 in zip(*axes):
         repeat = next(available_ix)
         x_ix[ax1] = repeat
diff --git a/opt_einsum/blas.py b/opt_einsum/blas.py
index 5ace3b6..14c12bc 100644
--- a/opt_einsum/blas.py
+++ b/opt_einsum/blas.py
@@ -6,6 +6,8 @@ import numpy as np
 
 from . import helpers
 
+__all__ = ["can_blas", "tensor_blas"]
+
 
 def can_blas(inputs, result, idx_removed, shapes=None):
     """
@@ -204,31 +206,27 @@ def tensor_blas(view_left, input_left, view_right, input_right, index_result, id
     # Matrix multiply
     # No transpose needed
     elif input_left[-rs:] == input_right[:rs]:
-        new_view = np.dot(view_left.reshape(dim_left, dim_removed),
-                          view_right.reshape(dim_removed, dim_right))
+        new_view = np.dot(view_left.reshape(dim_left, dim_removed), view_right.reshape(dim_removed, dim_right))
 
     # Transpose both
     elif input_left[:rs] == input_right[-rs:]:
-        new_view = np.dot(view_left.reshape(dim_removed, dim_left).T,
-                          view_right.reshape(dim_right, dim_removed).T)
+        new_view = np.dot(view_left.reshape(dim_removed, dim_left).T, view_right.reshape(dim_right, dim_removed).T)
 
     # Transpose right
     elif input_left[-rs:] == input_right[-rs:]:
-        new_view = np.dot(view_left.reshape(dim_left, dim_removed),
-                          view_right.reshape(dim_right, dim_removed).T)
+        new_view = np.dot(view_left.reshape(dim_left, dim_removed), view_right.reshape(dim_right, dim_removed).T)
 
     # Tranpose left
     elif input_left[:rs] == input_right[:rs]:
-        new_view = np.dot(view_left.reshape(dim_removed, dim_left).T,
-                          view_right.reshape(dim_removed, dim_right))
+        new_view = np.dot(view_left.reshape(dim_removed, dim_left).T, view_right.reshape(dim_removed, dim_right))
 
     # Conventional tensordot
     else:
         # Find indices to contract over
         left_pos, right_pos = (), ()
         for s in idx_removed:
-            left_pos += (input_left.find(s),)
-            right_pos += (input_right.find(s),)
+            left_pos += (input_left.find(s), )
+            right_pos += (input_right.find(s), )
         new_view = np.tensordot(view_left, view_right, axes=(left_pos, right_pos))
 
     # Make sure the resulting shape is correct
diff --git a/opt_einsum/compat.py b/opt_einsum/compat.py
index 38cf55d..a14846a 100644
--- a/opt_einsum/compat.py
+++ b/opt_einsum/compat.py
@@ -1,10 +1,12 @@
-# Python 2/3 compatability shim
+"""
+Python 2/3 compatability shim
+"""
 
 try:
     # Python 2
-    get_chr = unichr
-    strings = (str, type(get_chr(300)))
+    get_char = unichr
+    strings = (str, type(get_char(300)))
 except NameError:
     # Python 3
-    get_chr = chr
+    get_char = chr
     strings = str
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 01998dd..4d2f380 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -12,6 +12,8 @@ from . import parser
 from . import paths
 from . import sharing
 
+__all__ = ["contract_path", "contract", "format_const_einsum_str", "ContractExpression", "shape_only", "shape_only"]
+
 
 def contract_path(*operands, **kwargs):
     """
@@ -33,7 +35,7 @@ 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.
-        - 'auto' Use the optimal approach for small numbers of terms (currently
+        - 'auto' Use the optimal approach for a small number of terms (currently
           4 or less), else use the greedy approach.
 
     use_blas : bool
@@ -57,10 +59,10 @@ def contract_path(*operands, **kwargs):
     Examples
     --------
 
-    We can begin with a chain dot example. In this case it is optimal to
-    contract the b and c tensors reprsented by the first element of the path (1,
+    We can begin with a chain dot example. In this case, it is optimal to
+    contract the b and c tensors represented by the first element of the path (1,
     2). The resulting tensor is added to the end of the contraction and the
-    remaining contraction (0, 1) is then completed.
+    remaining contraction, ``(0, 1)``, is then executed.
 
     >>> a = np.random.rand(2, 2)
     >>> b = np.random.rand(2, 5)
@@ -271,7 +273,7 @@ def contract_path(*operands, **kwargs):
 
 @sharing.einsum_cache_wrap
 def _einsum(*operands, **kwargs):
-    """Base einsum, but with pre-parse for valid characters if string given.
+    """Base einsum, but with pre-parse for valid characters if a string is given.
     """
     fn = backends.get_func('einsum', kwargs.pop('backend', 'numpy'))
 
@@ -318,7 +320,7 @@ def contract(*operands, **kwargs):
     contract(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', use_blas=True, optimize=True, memory_limit=None, backend='numpy')
 
     Evaluates the Einstein summation convention on the operands. A drop in
-    replacment for NumPy's einsum function that optimizes the order of contraction
+    replacement for NumPy's einsum function that optimizes the order of contraction
     to reduce overall scaling at the cost of several intermediate arrays.
 
     Parameters
@@ -349,7 +351,7 @@ def contract(*operands, **kwargs):
           the number of terms in the contraction.
 
     memory_limit : int or None (default : None)
-        The upper limit of the size of tensor created, by default this will be
+        The upper limit of the size of tensor created, by default, this will be
         Give the upper bound of the largest intermediate tensor contract will build.
         By default (None) will size the ``memory_limit`` as the largest input tensor.
         Users can also specify ``-1`` to allow arbitrarily large tensors to be built.
@@ -366,7 +368,7 @@ def contract(*operands, **kwargs):
 
     Notes
     -----
-    This function should produce result identical to that of NumPy's einsum
+    This function should produce a result identical to that of NumPy's einsum
     function. The primary difference is ``contract`` will attempt to form
     intermediates which reduce the overall scaling of the given einsum contraction.
     By default the worst intermediate formed will be equal to that of the largest
@@ -374,9 +376,9 @@ def contract(*operands, **kwargs):
     provide arbitrarily large (1000 fold+) speed improvements.
 
     For contractions with just two tensors this function will attempt to use
-    NumPy's built in BLAS functionality to ensure that the given operation is
-    preformed in an optimal manner. When NumPy is linked to a threaded BLAS, potenital
-    speedsups are on the order of 20-100 for a six core machine.
+    NumPy's built-in BLAS functionality to ensure that the given operation is
+    preformed optimally. When NumPy is linked to a threaded BLAS, potential
+    speedups are on the order of 20-100 for a six core machine.
 
     Examples
     --------
@@ -579,11 +581,16 @@ class ContractExpression:
         """
         contraction_list = self._full_contraction_list if evaluate_constants else self.contraction_list
 
-        return _core_contract(list(arrays), contraction_list, out=out, backend=backend,
-                              evaluate_constants=evaluate_constants, **self.einsum_kwargs)
+        return _core_contract(
+            list(arrays),
+            contraction_list,
+            out=out,
+            backend=backend,
+            evaluate_constants=evaluate_constants,
+            **self.einsum_kwargs)
 
     def _contract_with_conversion(self, arrays, out, backend, evaluate_constants=False):
-        """Special contraction, i.e. contraction with a different backend
+        """Special contraction, i.e., contraction with a different backend
         but converting to and from that backend. Retrieves or generates a
         cached expression using ``arrays`` as templates, then calls it
         with ``arrays``.
@@ -641,7 +648,7 @@ class ContractExpression:
         try:
             # Check if the backend requires special preparation / calling
             #   but also ignore non-numpy arrays -> assume user wants same type back
-            if backend in backends.CONVERT_BACKENDS and any(isinstance(x, np.ndarray) for x in arrays):
+            if backends.has_backend(backend) and any(isinstance(x, np.ndarray) for x in arrays):
                 return self._contract_with_conversion(ops, out, backend, evaluate_constants=evaluate_constants)
 
             return self._contract(ops, out, backend, evaluate_constants=evaluate_constants)
@@ -679,8 +686,8 @@ def shape_only(shape):
 
 
 def contract_expression(subscripts, *shapes, **kwargs):
-    """Generate an reusable expression for a given contraction with
-    specific shapes, which can for example be cached.
+    """Generate a reusable expression for a given contraction with
+    specific shapes, which can, for example, be cached.
 
     Parameters
     ----------
@@ -692,8 +699,8 @@ def contract_expression(subscripts, *shapes, **kwargs):
         The indices of any constant arguments in ``shapes``, in which case the
         actual array should be supplied at that position rather than just a
         shape. If these are specified, then constant parts of the contraction
-        between calls will be reused. Additionally, if a gpu-enabled backend is
-        used for example, then the constant tensors will be kept on the gpu,
+        between calls will be reused. Additionally, if a GPU-enabled backend is
+        used for example, then the constant tensors will be kept on the GPU,
         minimizing transfers.
     kwargs :
         Passed on to ``contract_path`` or ``einsum``. See ``contract``.
@@ -714,7 +721,7 @@ def contract_expression(subscripts, *shapes, **kwargs):
     - 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.
-    - Constant operations will be computed upon first call with a particular
+    - Constant operations will be computed upon the first call with a particular
       backend, then subsequently reused.
 
     Examples
diff --git a/opt_einsum/helpers.py b/opt_einsum/helpers.py
index 9caa963..9efe5ea 100644
--- a/opt_einsum/helpers.py
+++ b/opt_einsum/helpers.py
@@ -4,9 +4,11 @@ Contains helper functions for opt_einsum testing scripts
 
 import numpy as np
 
-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)}
+__all__ = ["build_views", "compute_size_by_dict", "find_contraction", "flop_count"]
+
+_valid_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(_valid_chars, _sizes)}
 
 
 def build_views(string, dimension_dict=None):
@@ -18,7 +20,7 @@ def build_views(string, dimension_dict=None):
     tensor_list : list of str
         List of tensor strings to build
     dimension_dictionary : dictionary
-        Dictionary of index sizes
+        Dictionary of index _sizes
 
     Returns
     -------
@@ -34,7 +36,7 @@ def build_views(string, dimension_dict=None):
     """
 
     if dimension_dict is None:
-        dimension_dict = default_dim_dict
+        dimension_dict = _default_dim_dict
 
     views = []
     terms = string.split('->')[0].split(',')
@@ -54,7 +56,7 @@ def compute_size_by_dict(indices, idx_dict):
     indices : iterable
         Indices to base the product on.
     idx_dict : dictionary
-        Dictionary of index sizes
+        Dictionary of index _sizes
 
     Returns
     -------
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index adfc19f..2e9330e 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -10,17 +10,39 @@ import numpy as np
 
 from . import compat
 
-einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
+__all__ = [
+    "is_valid_einsum_char", "has_valid_einsum_chars_only", "get_symbol", "gen_unused_symbols",
+    "convert_to_valid_einsum_chars", "alpha_canonicalize", "find_output_str", "find_output_shape",
+    "possibly_convert_to_numpy", "parse_einsum_input"
+]
+
+_einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 
 
 def is_valid_einsum_char(x):
     """Check if the character ``x`` is valid for numpy einsum.
+
+    Examples
+    --------
+    >>> is_valid_einsum_char("a")
+    True
+
+    >>> is_valid_einsum_char("Ǵ")
+    False
     """
-    return (x in einsum_symbols_base) or (x in ',->.')
+    return (x in _einsum_symbols_base) or (x in ',->.')
 
 
 def has_valid_einsum_chars_only(einsum_str):
     """Check if ``einsum_str`` contains only valid characters for numpy einsum.
+
+    Examples
+    --------
+    >>> has_valid_einsum_chars_only("abAZ")
+    True
+
+    >>> has_valid_einsum_chars_only("Över")
+    False
     """
     return all(map(is_valid_einsum_char, einsum_str))
 
@@ -34,19 +56,24 @@ def get_symbol(i):
     >>> get_symbol(2)
     'c'
 
-    >>> oe.get_symbol(200)
+    >>> get_symbol(200)
     'Ŕ'
 
-    >>> oe.get_symbol(20000)
+    >>> get_symbol(20000)
     '京'
     """
     if i < 52:
-        return einsum_symbols_base[i]
-    return compat.get_chr(i + 140)
+        return _einsum_symbols_base[i]
+    return compat.get_char(i + 140)
 
 
 def gen_unused_symbols(used, n):
     """Generate ``n`` symbols that are not already in ``used``.
+
+    Examples
+    --------
+    >>> list(oe.parser.gen_unused_symbols("abd", 2))
+    ['c', 'e']
     """
     i = cnt = 0
     while cnt < n:
@@ -62,6 +89,11 @@ def convert_to_valid_einsum_chars(einsum_str):
     """Convert the str ``einsum_str`` to contain only the alphabetic characters
     valid for numpy einsum. If there are too many symbols, let the backend
     throw an error.
+
+    Examples
+    --------
+    >>> oe.parser.convert_to_valid_einsum_chars("Ĥěļļö")
+    'cbdda'
     """
     symbols = sorted(set(einsum_str) - set(',->'))
     replacer = {x: get_symbol(i) for i, x in enumerate(symbols)}
@@ -70,6 +102,14 @@ def convert_to_valid_einsum_chars(einsum_str):
 
 def alpha_canonicalize(equation):
     """Alpha convert an equation in an order-independent canonical way.
+
+    Examples
+    --------
+    >>> oe.parser.alpha_canonicalize("dcba")
+    'abcd'
+
+    >>> oe.parser.alpha_canonicalize("Ĥěļļö")
+    'abccd'
     """
     rename = OrderedDict()
     for name in equation:
@@ -81,7 +121,19 @@ def alpha_canonicalize(equation):
 
 
 def find_output_str(subscripts):
-    """Find the output string for the inputs ``susbcripts``.
+    """
+    Find the output string for the inputs ``subscripts`` under canonical einstein summation rules. That is, repeated indices are summed over by default.
+
+    Examples
+    --------
+    >>> oe.parser.find_output_str("ab,bc")
+    'ac'
+
+    >>> oe.parser.find_output_str("a,b")
+    'ab'
+
+    >>> oe.parser.find_output_str("a,a,b,b")
+    ''
     """
     tmp_subscripts = subscripts.replace(",", "")
     return "".join(s for s in sorted(set(tmp_subscripts)) if tmp_subscripts.count(s) == 1)
@@ -90,16 +142,45 @@ def find_output_str(subscripts):
 def find_output_shape(inputs, shapes, output):
     """Find the output shape for given inputs, shapes and output string, taking
     into account broadcasting.
+
+    Examples
+    --------
+    >>> oe.parser.find_output_shape(["ab", "bc"], [(2, 3), (3, 4)], "ac")
+    (2, 4)
+
+    # Broadcasting is accounted for
+    >>> oe.parser.find_output_shape(["a", "a"], [(4, ), (1, )], "a")
+    (4,)
     """
     return tuple(
-        max(shape[loc] for shape, loc in zip(shapes, [x.find(c) for x in inputs]) if loc >= 0)
-        for c in output
-    )
+        max(shape[loc] for shape, loc in zip(shapes, [x.find(c) for x in inputs]) if loc >= 0) for c in output)
 
 
 def possibly_convert_to_numpy(x):
     """Convert things without a 'shape' to ndarrays, but leave everything else.
+
+    Examples
+    --------
+    >>> oe.parser.possibly_convert_to_numpy(5)
+    array(5)
+
+    >>> oe.parser.possibly_convert_to_numpy([5, 3])
+    array([5, 3])
+
+    >>> oe.parser.possibly_convert_to_numpy(np.array([5, 3]))
+    array([5, 3])
+
+    # Any class with a shape is passed through
+    >>> class Shape:
+    ...     def __init__(self, shape):
+    ...         self.shape = shape
+    ...
+
+    >>> myshape = Shape((5, 5))
+    >>> oe.parser.possibly_convert_to_numpy(myshape)
+    <__main__.Shape object at 0x10f850710>
     """
+
     if not hasattr(x, 'shape'):
         return np.asanyarray(x)
     else:
diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py
index 0e00efe..8dc6009 100644
--- a/opt_einsum/paths.py
+++ b/opt_einsum/paths.py
@@ -5,6 +5,7 @@ import itertools
 
 from . import helpers
 
+__all__ = ["optimal", "greedy"] 
 
 def optimal(input_sets, output_set, idx_dict, memory_limit):
     """
diff --git a/opt_einsum/sharing.py b/opt_einsum/sharing.py
index 038309f..a982165 100644
--- a/opt_einsum/sharing.py
+++ b/opt_einsum/sharing.py
@@ -1,20 +1,28 @@
-# Copyright (c) 2018 Uber Technologies
+"""
+A module for sharing intermediates between contractions.
+
+Copyright (c) 2018 Uber Technologies
+"""
 
 import contextlib
 import functools
 import numbers
-from collections import Counter, defaultdict
 import threading
+from collections import Counter, defaultdict
 
 from .parser import alpha_canonicalize, parse_einsum_input
 
+__all__ = [
+    "currently_sharing", "get_sharing_cache", "shared_intermediates", "count_cached_ops", "transpose_cache_wrap",
+    "einsum_cache_wrap", "to_backend_cache_wrap"
+]
 
 _SHARING_STACK = defaultdict(list)
 
-
 try:
     get_thread_id = threading.get_ident
 except AttributeError:
+
     def get_thread_id():
         return threading.current_thread().ident
 
@@ -157,11 +165,15 @@ def einsum_cache_wrap(einsum):
         equation = args[0]
         inputs, output, operands = parse_einsum_input(args)
         inputs = inputs.split(',')
+
         _save_tensors(*operands)
+
+        # Build canonical key
         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(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 fdfacc7..64fb275 100644
--- a/opt_einsum/tests/test_contract.py
+++ b/opt_einsum/tests/test_contract.py
@@ -124,7 +124,7 @@ def test_compare_greek(string):
     ein = contract(string, *views, optimize=False, use_blas=False)
 
     # convert to greek
-    string = ''.join(compat.get_chr(ord(c) + 848) if c not in ',->.' else c for c in string)
+    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)
@@ -152,7 +152,7 @@ def test_compare_blas_greek(string):
     ein = contract(string, *views, optimize=False)
 
     # convert to greek
-    string = ''.join(compat.get_chr(ord(c) + 848) if c not in ',->.' else c for c in string)
+    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)
@@ -163,7 +163,7 @@ def test_compare_blas_greek(string):
 
 def test_some_non_alphabet_maintains_order():
     # 'c beta a' should automatically go to -> 'a c beta'
-    string = 'c' + compat.get_chr(ord('b') + 848) + 'a'
+    string = 'c' + compat.get_char(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)