Repository: dgasmith/opt_einsum
Pull request #17 (2018-03-15 13:15:41)

non-numpy (e.g. GPU) backend support
This adds non-numpy backend support to ``opt_einsum`` in two flavors: 

### General duck-typed support

For arbitrary ndarray libraries, i.e. if the arrays supplied have got a ``.shape`` attribute and opt_einsum can find ``{backend}.tensordot`` etc then it should work, returning whatever type the library uses. So for example 

```python
contract("abc,bcd,def->aef", *dask_arrays, backend='dask')
```

 will return a ``dask.Array`` without any explicit logic. I've added tests for [dask](https://github.com/dask/dask/issues/732) and also [sparse](https://github.com/pydata/sparse/issues/31).

Edit: I should add that extra functionality like taking diagonals etc is still reliant on a ``einsum`` implementation. Not sure about broadcasting status etc.

### Special conversion support

To support GPU operations on numpy arrays, here, on first call the arrays are converted to the correct placeholder, an expression generated using these, and then this expression can be called with normal numpy arrays, returning numpy arrays also. E.g.:

```python
expr = contract_expression("abc,bcd,def->aef", *shapes)
expr(*numpy_arrays, backend='theano')
```

Explicit tested support here is for ``tensorflow``, ``theano`` and ``cupy`` - GPU libraries where it makes sense to convert to and from.


### Adding more backends

``contract.py`` doesn't have any explicit reference to the various backends, thats all contained in ``backends.py`` now. So it should be very easy to add new backends. E.g. if ``tblis`` were to add ``tensordot``, it would fall under the duck-typing category above and I think no changes would actually be needed at all.

Obviously any recommendations welcome! Getting the right backend function and all that is cached so I don't think this has added any more than a few 100ns overhead.

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


jcmgray: I've update things so that backends without ``einsum`` can do as many types of contraction as possible but will still error more clearly if they need einsum - this needed a little bit of a refactor of ``can_blas``.

There's also a built version of the added docs [here](http://opt-einsum-jcmgray.readthedocs.io/en/latest/backends.html) if you want to check that out.

By the way, I think ``greedy`` now scales as N^2, (N^2 initial cost plus N steps of checking <N contractions), was thinking I might just update those docs in this PR as well?
dgasmith: Nice, I think the docs look good. Would it be too much to give an example for each backend along with labels so that users can skip right to that part they may be interested in?

For refactoring `can_blas` it might be nice to screen based on the type of BLAS operation it is. Perhaps like the following:
 - `False` or 0 - No BLAS possible
 - 1, 2, 3 - Standard BLAS level
 - 4 - BLAS operations that cannot be written as standard level N BLAS such as ('ijkl,jlab->ikab')

Sure please do update the greedy docs.
dgasmith: Another optimization to consider and may do later is `numexpr`. It would be quite good for typical einsum like reductions. 
jcmgray: Sorry the delay! I think the docs recommendation is a good one for ease of access- will do that soon. Some other thoughts:

The BLAS recommendation also sounds nice - but I might leave for another time since here I think the requirement here is only 'should-blas', 'can-blas' or 'no-blas'.

``numba`` has just added support for transposing arbitrary ndarrays, so that will be an option too (would have to write a ``tensordot`` function, but it should be a bit faster than ``numpy``.

In general - and relating to ``numexpr`` - if one could perform the remaining ``einsum`` operations like axis reductions but without einsum, then you'd have a pretty complete replacement for einsum for any tensor library, which would just need some base requirement functions. i.e. something like:

```python
if can_blas:
    ...
    _tensordot(x, y, ...)
    ...
    _transpose(...)
    ...
else:
    _axes_reduction(x, y, ...)
```
where _axes_reduction just needs (I'm guessing) ``diag``, ``sum`` etc.
Are axis reductions the only non-tensordot-able contractions?
dgasmith: I think axis reductions and Hadamard products are the only non-tensordot-able products (perhaps by definition?). That was my thought as well, we can fall back to numexpr or numba in certain cases and let the other backend handle the heavy lifting. My only concern with this is moving data on and off GPU's is quite expensive and it would be preferable to just keep the data on the GPU's if possible. I would say lets make an issue of this and get this PR in as is before it gets too large.

For the BLAS recommendation can you do a halfway change as I'm not very comfortable with `True`, `False`, `None` logic chains.
jcmgray: I think if ``_axes_reduction`` would take a ``backend`` argument then it could just get its functions in the same way as ``_tensordot`` and GPU arrays would stay on the GPU etc. But anyway, agree it is a problem for a separate 'einsum without einsum' issue.

And sure, I'll make some changes to ``can_blas`` to make the three cases a bit more explicit, it's definitely a bit arcane at the moment.
dgasmith: Agreed all around. Once the `can_blas` is committed I will pull this in. Also, the other PR will make codecov go away.

Patch:
diff --git a/.travis.yml b/.travis.yml
index 1e4eaa8..57069df 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -10,6 +10,7 @@ matrix:
     - python: 2.7
     - python: 3.4
     - python: 3.5
+    - python: 3.6
 
 before_install:
   - uname -a
@@ -21,6 +22,10 @@ before_install:
   - pip install pytest pytest-cov
   - pip install numpy
   - pip install codecov
+  - pip install tensorflow
+  - pip install theano
+  - pip install dask[array]
+  - pip install sparse
   - pip install -e .
 
 script:
diff --git a/README.md b/README.md
index 059b974..a6c5eeb 100644
--- a/README.md
+++ b/README.md
@@ -234,7 +234,7 @@ The final contraction cost is computed and we choose the second path from the li
 
 ## Finding the opportunistic path
 
-Another way to find a path is to choose the best pair to contract at every iteration so that the formula scales like N^3. 
+Another way to find a path is to choose the best pair to contract at every iteration so that the formula scales like N^2. 
 The "best" contraction pair is currently determined by the smallest of the tuple ``(-removed_size, cost)`` where ``removed_size`` is the size of the contracted tensors minus the size of the tensor created and ``cost`` is the cost of the contraction.
 Basically, we want to remove the largest dimensions at the least cost.
 To prevent large outer products the results are sieved by the amount of memory available.
@@ -242,6 +242,10 @@ Overall, this turns out to work extremely well and is only slower than the optim
 To me, while still not perfect, it represents a "good enough" algorithm for general production.
 It is fast enough that at worst case the overhead penalty is approximately 20 microseconds and is much faster for every other einsum test case I can build or generate randomly.
 
+Since :mod:`opt_einsum` can handle up to a thousand unique indices (which could be increased),
+this greedy approach is also efficient enough to find the contraction path for expressions with hundreds of tensors in less than a second.
+
+
 ## Testing
 
 Testing this function thoroughly is absolutely crucial; the testing scripts do required python pandas in addition to numpy. Testing is broken down into several tasks:
diff --git a/docs/requirements.yml b/docs/requirements.yml
new file mode 100644
index 0000000..803de5d
--- /dev/null
+++ b/docs/requirements.yml
@@ -0,0 +1,8 @@
+name: opt-einsum-docs
+channels:
+    - defaults
+dependencies:
+    - python=3
+    - numpy
+    - sphinx
+    - sphinx_rtd_theme
diff --git a/docs/source/backends.rst b/docs/source/backends.rst
new file mode 100644
index 0000000..f424158
--- /dev/null
+++ b/docs/source/backends.rst
@@ -0,0 +1,120 @@
+======================
+Backends & GPU Support
+======================
+
+A brief overview of libraries that have been tested:
+
+    - `tensorflow <https://www.tensorflow.org/>`_: compiled tensor expressions
+      that can run on GPU.
+    - `theano <http://deeplearning.net/software/theano/>`_: compiled tensor
+      expressions that can run on GPU.
+    - `cupy <https://cupy.chainer.org/>`_: numpy-like api for GPU tensors.
+    - `dask <https://dask.pydata.org/>`_: larger-than-memory tensor
+      computations, distributed scheduling, and potential reuse of
+      intermediaries.
+    - `sparse <https://sparse.pydata.org/>`_: sparse tensors.
+
+``opt_einsum`` is quite agnostic to the type of n-dimensional arrays (tensors)
+it uses, since finding the contraction path only relies on getting the shape
+attribute of each array supplied.
+It can perform the underlying tensor contractions with various
+libraries. In fact, any library that provides a :func:`~numpy.tensordot` and
+:func:`~numpy.transpose` implementation can perform most normal contractions.
+While more special functionality such as taking diagonals is reliant on a
+:func:`~numpy.einsum` implementation.
+
+
+
+General backend for any type of array
+-------------------------------------
+
+This 'duck-typing' support just requires specifying the correct ``backend``
+argument for the type of arrays supplied when calling
+:func:`~opt_einsum.contract`. For example with ``dask``:
+
+.. code-block:: python
+
+    >>> import opt_einsum as oe
+    >>> import dask.array as da
+    >>> shapes = (3, 200), (200, 300), (300, 4)
+    >>> dxs = [da.random.normal(0, 1, shp, chunks=(100, 100)) for shp in shapes]
+    >>> dxs
+    [dask.array<da.random.normal, shape=(3, 200), dtype=float64, chunksize=(3, 100)>,
+     dask.array<da.random.normal, shape=(200, 300), dtype=float64, chunksize=(100, 100)>,
+     dask.array<da.random.normal, shape=(300, 4), dtype=float64, chunksize=(100, 4)>]
+
+
+    >>> dy = oe.contract("ab,bc,cd", *dxs, backend='dask')
+    >>> dy
+    dask.array<transpose, shape=(3, 4), dtype=float64, chunksize=(3, 4)>
+
+    >>> dy.compute()
+    array([[ 470.71404665,    2.44931372,  -28.47577265,  424.37716615],
+           [  64.38328345, -287.40753131,  144.46515642,  324.88169821],
+           [-142.07153553, -180.41739259,  125.0973783 , -239.16754541]])
+
+
+In this case, dask arrays in = dask array out, since dask arrays have a shape
+attribute, and ``opt_einsum`` can find ``dask.array.tensordot`` and
+``dask.array.transpose``.
+
+
+
+Special Backends for numpy arrays - GPU etc.
+----------------------------------------------
+
+A special case is if you want to supply numpy arrays and get numpy arrays back,
+but use a different backend, such as performing a contraction on a GPU.
+Unless the specified backend works on numpy arrays this requires converting to
+and from the backend array type. Currently ``opt_einsum`` can handle this
+automatically for:
+
+    - ``tensorflow``
+    - ``theano``
+    - ``cupy``
+
+which all offer GPU support. Since ``tensorflow`` and ``theano`` both require
+compiling the expression, this functionality is encapsulated in generating a
+:class:`~opt_einsum.ContractExpression` using
+:func:`~opt_einsum.contract_expression`, which can then be called using numpy
+arrays whilst specifiying ``backend='tensorflow'`` etc.
+
+For example with **theano**:
+
+.. code-block:: python
+
+    >>> import opt_einsum as oe
+    >>> shapes = (3, 200), (200, 300), (300, 4)
+    >>> expr = oe.contract_expression("ab,bc,cd", *shapes)
+    >>> expr
+    ContractExpression('ab,bc,cd')
+
+    >>> import numpy as np
+    >>> # GPU advantage mainly for low precision numbers
+    >>> xs = xs = [np.random.randn(*shp).astype(np.float32) for shp in shapes]
+    >>> expr(*xs, backend='theano')  # might see some fluff on first run
+    ...
+    array([[ 129.28352  , -128.00702  , -164.62917  , -335.11682  ],
+           [-462.52344  , -121.12657  ,  -67.847626 ,  624.5457   ],
+           [   5.2838974,   36.441578 ,   81.62851  ,  703.1576   ]],
+          dtype=float32)
+
+To run the expression with **tensorflow**, you need to register a default
+session:
+
+.. code-block:: python
+
+    >>> import tensorflow as tf
+    >>> sess = tf.Session()  # might see some fluff
+    ...
+
+    >>> with sess.as_default(): out = expr(*xs, backend='tensorflow')
+    >>> out
+    array([[ 129.28357  , -128.00684  , -164.62903  , -335.1167   ],
+           [-462.52362  , -121.12659  ,  -67.84769  ,  624.5455   ],
+           [   5.2839584,   36.44155  ,   81.62852  ,  703.15784  ]],
+          dtype=float32)
+
+Note that one could still supply this expression with, for example, a
+``tensorflow.placeholder`` using ``backend='tensorflow'``, and then no
+conversion would take place, instead you'd get a ``tensorflow.Tensor`` back.
diff --git a/docs/source/greedy_path.rst b/docs/source/greedy_path.rst
index 6a7dc47..44ff956 100644
--- a/docs/source/greedy_path.rst
+++ b/docs/source/greedy_path.rst
@@ -2,10 +2,13 @@
 The Greedy Path
 ===============
 
-Another way to find a path is to choose the best pair to contract at every iteration so that the formula scales like N^3 - functionality provided by :func:`~opt_einsum.paths.greedy`.
+Another way to find a path is to choose the best pair to contract at every iteration so that the formula scales like N^2 - functionality provided by :func:`~opt_einsum.paths.greedy`.
 The "best" contraction pair is currently determined by the smallest of the tuple ``(-removed_size, cost)`` where ``removed_size`` is the size of the contracted tensors minus the size of the tensor created and ``cost`` is the cost of the contraction.
 Basically, we want to remove the largest dimensions at the least cost.
 To prevent large outer products the results are sieved by the amount of memory available.
 Overall, this turns out to work extremely well and is only slower than the optimal path in several cases, and even then only by a factor of 2-4 while only taking 1 millisecond for terms of length 10.
 To me, while still not perfect, it represents a "good enough" algorithm for general production.
 It is fast enough that at worst case the overhead penalty is approximately 20 microseconds and is much faster for every other einsum test case I can build or generate randomly.
+
+Since :mod:`opt_einsum` can handle up to a thousand unique indices (which could be increased),
+this greedy approach is also efficient enough to find the contraction path for expressions with hundreds of tensors in less than a second.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index aa089ea..bbc33d6 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -95,6 +95,7 @@ We can then view more details about the optimized contraction order:
 
    install
    examples
+   backends
 
 
 .. toctree::
diff --git a/opt_einsum/backends.py b/opt_einsum/backends.py
new file mode 100644
index 0000000..47fcbae
--- /dev/null
+++ b/opt_einsum/backends.py
@@ -0,0 +1,144 @@
+"""
+Performing array conversions and operations with various backends.
+"""
+
+import numpy
+import importlib
+
+
+# known non top-level imports
+_aliases = {'dask': 'dask.array', 'theano': 'theano.tensor'}
+
+
+def _import_func(func, backend):
+    """Try and import ``{backend}.{func}``, raise an error if library is
+    installed but can't find that particular function.
+    """
+    try:
+        lib = importlib.import_module(_aliases.get(backend, backend))
+        return getattr(lib, func)
+    except AttributeError:
+        raise AttributeError("{} doesn't seem to provide the function {}".format(backend, func))
+
+
+# manually cache functions as python2 doesn't support functools.lru_cache
+#     other libs will be added to this if needed, but pre-populate with numpy
+_cached_funcs = {
+    ('tensordot', 'numpy'): numpy.tensordot,
+    ('transpose', 'numpy'): numpy.transpose,
+    ('einsum', 'numpy'): numpy.einsum,
+}
+
+
+def get_func(func, backend='numpy'):
+    """Return ``{backend}.{func}``, e.g. ``numpy.einsum``, cache result.
+    """
+    try:
+        return _cached_funcs[func, backend]
+    except KeyError:
+        fn = _import_func(func, backend)
+        _cached_funcs[func, backend] = fn
+        return fn
+
+
+# mark libs with einsum, else try to use tensordot/tranpose as much as possible
+_has_einsum = {}
+
+
+def has_einsum(backend):
+    """Check if ``{backend}.einsum`` exists, cache result for performance.
+    """
+    try:
+        return _has_einsum[backend]
+    except KeyError:
+        try:
+            get_func('einsum', backend)
+            _has_einsum[backend] = True
+        except AttributeError:
+            _has_einsum[backend] = False
+
+        return _has_einsum[backend]
+
+
+# Tensorflow
+
+def convert_arrays_to_tensorflow(arrays):
+    """Convert numpy arrays to ``tensorflow.placeholder`` instances.
+    """
+    import tensorflow
+    return [tensorflow.placeholder(x.dtype, x.shape) for x in arrays]
+
+
+def build_tensorflow_expression(arrays, expr):
+    """Build a tensorflow function based on ``arrays`` and ``expr``.
+    """
+    import tensorflow
+    placeholders = convert_arrays_to_tensorflow(arrays)
+    graph = expr._normal_contract(placeholders, backend='tensorflow')
+
+    def tensorflow_contract(*arrays):
+        session = tensorflow.get_default_session()
+        return session.run(graph, feed_dict=dict(zip(placeholders, arrays)))
+
+    return tensorflow_contract
+
+
+# Theano
+
+def convert_arrays_to_theano(arrays):
+    """Convert numpy arrays to ``theano.tensor.TensorType`` instances.
+    """
+    import theano
+    return [theano.tensor.TensorType(dtype=x.dtype, broadcastable=[False] * len(x.shape))() for x in arrays]
+
+
+def build_theano_expression(arrays, expr):
+    """Build a theano function based on ``arrays`` and ``expr``.
+    """
+    import theano
+    in_vars = convert_arrays_to_theano(arrays)
+    out_var = expr._normal_contract(in_vars, backend='theano')
+    graph = theano.function(in_vars, out_var)
+
+    def theano_contract(*arrays):
+        return graph(*arrays)
+
+    return theano_contract
+
+
+# Cupy
+
+def convert_arrays_to_cupy(arrays):  # pragma: no cover
+    """Convert numpy arrays to ``cupy.ndarray`` instances.
+    """
+    import cupy
+    return [cupy.asarray(x) for x in arrays]
+
+
+def build_cupy_expression(_, expr):  # pragma: no cover
+    """Build a cupy function based on ``arrays`` and ``expr``.
+    """
+    import cupy
+
+    def cupy_contract(*arrays):
+        cupy_arrays = convert_arrays_to_cupy(arrays)
+        cupy_out = expr._normal_contract(cupy_arrays, backend='cupy')
+        return cupy.asnumpy(cupy_out)
+
+    return cupy_contract
+
+
+# Dispatch to correct expression backend
+#    these are the backends which support explicit to-and-from numpy conversion
+CONVERT_BACKENDS = {
+    'tensorflow': build_tensorflow_expression,
+    'theano': build_theano_expression,
+    'cupy': build_cupy_expression,
+}
+
+
+def build_expression(backend, arrays, expr):
+    """Build an expression, based on ``expr`` and initial arrays ``arrays``,
+    that evaluates using backend ``backend``.
+    """
+    return CONVERT_BACKENDS[backend](arrays, expr)
diff --git a/opt_einsum/blas.py b/opt_einsum/blas.py
index 93f878a..b8170ce 100644
--- a/opt_einsum/blas.py
+++ b/opt_einsum/blas.py
@@ -29,7 +29,9 @@ def can_blas(inputs, result, idx_removed):
     Notes
     -----
     We assume several operations are not efficient such as a transposed
-    DDOT, therefore 'ijk,jki->' would return False.
+    DDOT, therefore 'ijk,jki->' should prefer einsum. These return the blas
+    type appended with "/EINSUM" to differentiate when they can still be done
+    with tensordot if required, e.g. when a backend has no einsum.
 
     Examples
     --------
@@ -39,45 +41,48 @@ def can_blas(inputs, result, idx_removed):
     >>> _can_blas(['ijj', 'jk'], 'ik', set('j'))
     False
 
+    >>> _can_blas(['ab', 'cd'], 'abcd', set())
+    'OUTER/EINSUM'
     """
-
-    # Gotta remove indices
-    if len(idx_removed) == 0:
-        return False
-
     # Can only do two
     if len(inputs) != 2:
         return False
 
-    # Make sure there is overlap
-    if len(set(inputs[0]) & set(inputs[1])) == 0:
-        return False
+    input_left, input_right = inputs
+
+    for c in set(input_left + input_right):
+        # can't deal with repeated indices on same input or more than 2 total
+        nl, nr = input_left.count(c), input_right.count(c)
+        if (nl > 1) or (nr > 1) or (nl + nr > 2):
+            return False
+
+        # can't do implicit summation or dimension collapse e.g.
+        #     "ab,bc->c" (implicitly sum over 'a')
+        #     "ab,ca->ca" (take diagonal of 'a')
+        if nl + nr - 1 == int(c in result):
+            return False
+
+    # Prefer einsum if not removing indices
+    #     (N.B. tensordot outer faster for large arrays?)
+    if len(idx_removed) == 0:
+        return 'OUTER/EINSUM'
 
     # Build a few temporaries
     sets = [set(x) for x in inputs]
     keep_left = sets[0] - idx_removed
     keep_right = sets[1] - idx_removed
-    input_left = inputs[0]
-    input_right = inputs[1]
     rs = len(idx_removed)
 
-    if any(len(l) != len(s) for l, s in zip(inputs, sets)):
-        return False
-
-    # Cannot handle partial inner
-    if len(keep_left & keep_right):
-        return False
-
     # DDOT
-    elif inputs[0] == inputs[1]:
+    if inputs[0] == inputs[1]:
         return 'DOT'
 
-    # DDOT doesnt make sense if you have to tranpose
+    # DDOT doesnt make sense if you have to tranpose - prefer einsum
     elif sets[0] == sets[1]:
-        return False
+        return 'DOT/EINSUM'
 
     # GEMM no transpose
-    elif input_left[-rs:] == input_right[:rs]:
+    if input_left[-rs:] == input_right[:rs]:
         return 'GEMM'
 
     # GEMM transpose both
@@ -94,7 +99,7 @@ def can_blas(inputs, result, idx_removed):
 
     # Einsum is faster than vectordot if we have to copy
     elif (len(keep_left) == 0) or (len(keep_right) == 0):
-        return False
+        return 'GEMV/EINSUM'
 
     # Conventional tensordot
     else:
@@ -149,7 +154,6 @@ def tensor_blas(view_left, input_left, view_right, input_right, index_result, id
     idx_removed = set(idx_removed)
     keep_left = set(input_left) - idx_removed
     keep_right = set(input_right) - idx_removed
-    #print(input_left, input_right, idx_removed)
 
     # We trust this must be called correctly
     dimension_dict = {}
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 08b7ca1..98c5da5 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -8,6 +8,7 @@ from . import blas
 from . import helpers
 from . import paths
 from . import parser
+from . import backends
 
 
 def contract_path(*operands, **kwargs):
@@ -242,23 +243,25 @@ def contract_path(*operands, **kwargs):
     path_print += "   Theoretical speedup:  %3.3f\n" % (naive_cost / float(opt_cost))
     path_print += "  Largest intermediate:  %.3e elements\n" % max(size_list)
     path_print += "-" * 80 + "\n"
-    path_print += "%6s %6s %24s %40s\n" % header
+    path_print += "%6s %11s %22s %37s\n" % header
     path_print += "-" * 80
 
     for n, contraction in enumerate(contraction_list):
         inds, idx_rm, einsum_str, remaining, do_blas = contraction
         remaining_str = ",".join(remaining) + "->" + output_subscript
         path_run = (scale_list[n], do_blas, einsum_str, remaining_str)
-        path_print += "\n%4d %9s %24s %40s" % path_run
+        path_print += "\n%4d %14s %22s %37s" % path_run
 
     return path, path_print
 
 
-def _np_einsum(*operands, **kwargs):
-    """Numpy einsum, but with pre-parse for valid characters if string given.
+def _einsum(*operands, **kwargs):
+    """Base einsum, but with pre-parse for valid characters if string given.
     """
+    fn = backends.get_func('einsum', kwargs.pop('backend', 'numpy'))
+
     if not isinstance(operands[0], str):
-        return np.einsum(*operands, **kwargs)
+        return fn(*operands, **kwargs)
 
     einsum_str, operands = operands[0], operands[1:]
 
@@ -271,14 +274,32 @@ def _np_einsum(*operands, **kwargs):
 
         einsum_str = parser.convert_to_valid_einsum_chars(einsum_str)
 
-    return np.einsum(einsum_str, *operands, **kwargs)
+    return fn(einsum_str, *operands, **kwargs)
+
+
+def _transpose(x, axes, backend='numpy'):
+    """Base transpose.
+    """
+    try:
+        return x.transpose(axes)
+    except AttributeError:
+        # some libraries don't implement method version
+        fn = backends.get_func('transpose', backend)
+        return fn(x, axes)
+
+
+def _tensordot(x, y, axes, backend='numpy'):
+    """Base tensordot.
+    """
+    fn = backends.get_func('tensordot', backend)
+    return fn(x, y, axes=axes)
 
 
 # Rewrite einsum to handle different cases
 def contract(*operands, **kwargs):
     """
-    contract(subscripts, *operands, out=None, dtype=None, order='K',
-           casting='safe', use_blas=True, optimize=True, memory_limit=None)
+    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
@@ -316,6 +337,11 @@ def contract(*operands, **kwargs):
         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.
+    backend : str, optional (default: ``numpy``)
+        Which library to use to perform the required ``tensordot``, ``transpose``
+        and ``einsum`` calls. Should match the types of arrays supplied, See
+        :func:`contract_expression` for generating expressions which convert
+        numpy arrays to and from the backend library automatically.
 
     Returns
     -------
@@ -351,11 +377,12 @@ def contract(*operands, **kwargs):
 
     # If no optimization, run pure einsum
     if optimize_arg is False:
-        return _np_einsum(*operands, **einsum_kwargs)
+        return _einsum(*operands, **einsum_kwargs)
 
     # Grab non-einsum kwargs
     use_blas = kwargs.pop('use_blas', True)
     memory_limit = kwargs.pop('memory_limit', None)
+    backend = kwargs.pop('backend', 'numpy')
     gen_expression = kwargs.pop('gen_expression', False)
 
     # Make sure remaining keywords are valid for einsum
@@ -374,10 +401,10 @@ def contract(*operands, **kwargs):
     if gen_expression:
         return ContractExpression(full_str, contraction_list, **einsum_kwargs)
 
-    return _core_contract(operands, contraction_list, **einsum_kwargs)
+    return _core_contract(operands, contraction_list, backend=backend, **einsum_kwargs)
 
 
-def _core_contract(operands, contraction_list, **einsum_kwargs):
+def _core_contract(operands, contraction_list, backend='numpy', **einsum_kwargs):
     """Inner loop used to perform an actual contraction given the output
     from a ``contract_path(..., einsum_call=True)`` call.
     """
@@ -386,6 +413,9 @@ def _core_contract(operands, contraction_list, **einsum_kwargs):
     out_array = einsum_kwargs.pop('out', None)
     specified_out = out_array is not None
 
+    # try and do as much as possible without einsum if not available
+    no_einsum = not backends.has_einsum(backend)
+
     # Start contraction loop
     for num, contraction in enumerate(contraction_list):
         inds, idx_rm, einsum_str, remaining, blas = contraction
@@ -396,8 +426,8 @@ def _core_contract(operands, contraction_list, **einsum_kwargs):
         # Do we need to deal with the output?
         handle_out = specified_out and ((num + 1) == len(contraction_list))
 
-        # Call tensordot
-        if blas:
+        # Call tensordot (check if should prefer einsum, but only if available)
+        if blas and ('EINSUM' not in blas or no_einsum):
 
             # Checks have already been handled
             input_str, results_index = einsum_str.split('->')
@@ -414,13 +444,16 @@ def _core_contract(operands, contraction_list, **einsum_kwargs):
                 right_pos.append(input_right.find(s))
 
             # Contract!
-            new_view = np.tensordot(*tmp_operands, axes=(tuple(left_pos), tuple(right_pos)))
+            new_view = _tensordot(*tmp_operands, axes=(tuple(left_pos), tuple(right_pos)), backend=backend)
 
             # Build a new view if needed
             if (tensor_result != results_index) or handle_out:
+
+                transpose = tuple(map(tensor_result.index, results_index))
+                new_view = _transpose(new_view, axes=transpose, backend=backend)
+
                 if handle_out:
-                    einsum_kwargs["out"] = out_array
-                new_view = _np_einsum(tensor_result + '->' + results_index, new_view, **einsum_kwargs)
+                    out_array[:] = new_view
 
         # Call einsum
         else:
@@ -429,7 +462,7 @@ def _core_contract(operands, contraction_list, **einsum_kwargs):
                 einsum_kwargs["out"] = out_array
 
             # Do the contraction
-            new_view = _np_einsum(einsum_str, *tmp_operands, **einsum_kwargs)
+            new_view = _einsum(einsum_str, *tmp_operands, backend=backend, **einsum_kwargs)
 
         # Append new items and derefernce what we can
         operands.append(new_view)
@@ -452,18 +485,64 @@ class ContractExpression:
         self.einsum_kwargs = einsum_kwargs
         self.num_args = len(contraction.split('->')[0].split(','))
 
+    def _normal_contract(self, arrays, out=None, backend='numpy'):
+        """The normal, core contraction.
+        """
+        return _core_contract(list(arrays), self.contraction_list, out=out,
+                              backend=backend, **self.einsum_kwargs)
+
+    def _convert_contract(self, arrays, out, backend):
+        """Special contraction, i.e. contraction with a different backend
+        but converting to and from that backend. Checks for
+        ``self._{backend}_contract``, generates it if is missing, then calls it
+        with ``arrays``.
+        """
+        convert_fn = "_{}_contract".format(backend)
+
+        if not hasattr(self, convert_fn):
+            setattr(self, convert_fn, backends.build_expression(backend, arrays, self))
+
+        result = getattr(self, convert_fn)(*arrays)
+
+        if out is not None:
+            out[()] = result
+            return out
+
+        return result
+
     def __call__(self, *arrays, **kwargs):
+        """Evaluate this expression with a set of arrays.
+
+        Parameters
+        ----------
+        arrays : seq of array
+            The arrays to supply as input to the expression.
+        out : array, optional (default: ``None``)
+            If specified, output the result into this array.
+        backend : str, optional  (default: ``numpy``)
+            Perform the contraction with this backend library. If numpy arrays
+            are supplied then try to convert them to and from the correct
+            backend array type.
+        """
+
         if len(arrays) != self.num_args:
             raise ValueError("This `ContractExpression` takes exactly %s array arguments "
                              "but received %s." % (self.num_args, len(arrays)))
 
+        backend = kwargs.pop('backend', 'numpy')
         out = kwargs.pop('out', None)
         if kwargs:
-            raise ValueError("The only valid keyword argument to a `ContractExpression` "
-                             "call is `out=`. Got: %s." % kwargs)
+            raise ValueError("The only valid keyword arguments to a `ContractExpression` "
+                             "call are `out=` or `backend=`. Got: %s." % kwargs)
 
         try:
-            return _core_contract(list(arrays), self.contraction_list, out=out, **self.einsum_kwargs)
+            # 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 isinstance(arrays[0], np.ndarray):
+                return self._convert_contract(arrays, out, backend)
+
+            return self._normal_contract(arrays, out, backend)
+
         except ValueError as err:
             original_msg = str(err.args) if err.args else ""
             msg = ("Internal error while evaluating `ContractExpression`. Note that few checks are performed"
@@ -510,13 +589,16 @@ def contract_expression(subscripts, *shapes, **kwargs):
     Returns
     -------
     expr : ContractExpression
-        Callable with signature ``expr(*arrays)`` where the array's shapes
-        should match ``shapes``.
+        Callable with signature ``expr(*arrays, out=None, backend='numpy')``
+        where the array's shapes should match ``shapes``.
 
     Notes
     -----
     - The `out` keyword argument should be supplied to the generated expression
       rather than this function.
+    - The `backend` keyword argument should also be supplied to the generated
+      expression. If numpy arrays are supplied, if possible they will be
+      converted to and back from the correct backend array type.
     - 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.
@@ -534,8 +616,10 @@ def contract_expression(subscripts, *shapes, **kwargs):
     if not kwargs.get('optimize', True):
         raise ValueError("Can only generate expressions for optimized contractions.")
 
-    if kwargs.get('out', None) is not None:
-        raise ValueError("`out` should only be specified when calling a `ContractExpression`, not when building it.")
+    for arg in ('out', 'backend'):
+        if kwargs.get(arg, None) is not None:
+            raise ValueError("'%s' should only be specified when calling a "
+                             "`ContractExpression`, not when building it." % arg)
 
     dummy_arrays = [_ShapeOnly(s) for s in shapes]
 
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index e973826..9d51c87 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -5,8 +5,8 @@ A functionally equivalent parser of the numpy.einsum input parser
 import numpy as np
 import sys
 
-einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
-einsum_symbols = str(einsum_symbols_base)
+einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyz'
+einsum_symbols = einsum_symbols_base + 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
 
 # boost the number of symbols using unicode if python3
 if sys.version_info[0] >= 3:
diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py
index 7df0410..0e6759b 100644
--- a/opt_einsum/paths.py
+++ b/opt_einsum/paths.py
@@ -11,7 +11,7 @@ def optimal(input_sets, output_set, idx_dict, memory_limit):
     on ``memory_limit`` and returns the lowest cost path. This algorithm
     scales factorial with respect to the elements in the list ``input_sets``.
 
-    Paramaters
+    Parameters
     ----------
     input_sets : list
         List of sets that represent the lhs side of the einsum subscript
@@ -181,13 +181,14 @@ def greedy(input_sets, output_set, idx_dict, memory_limit):
     """
     Finds the path by contracting the best pair until the input list is
     exhausted. The best pair is found by minimizing the tuple
-    ``(-prod(indices_removed), cost)``.  What this amounts to is prioritizing
-    matrix multiplication or inner product operations, then Hadamard like
+    ``(-removed_size, cost)``.  What this amounts to is prioritizing
+    inner product operations, matrix multiplication, then Hadamard like
     operations, and finally outer operations. Outer products are limited by
-    ``memory_limit``. This algorithm scales cubically with respect to the number of
-    elements in the list ``input_sets``.
+    ``memory_limit`` and are ignored until no other operations are
+    available. This algorithm scales quadratically with respect to the
+    number of elements in the list ``input_sets``.
 
-    Paramaters
+    Parameters
     ----------
     input_sets : list
         List of sets that represent the lhs side of the einsum subscript
diff --git a/readthedocs.yml b/readthedocs.yml
new file mode 100644
index 0000000..1bc7609
--- /dev/null
+++ b/readthedocs.yml
@@ -0,0 +1,5 @@
+conda:
+    file: docs/requirements.yml
+python:
+    version: 3
+    setup_py_install: true


Test patch:
diff --git a/tests/test_backends.py b/tests/test_backends.py
new file mode 100644
index 0000000..a8c8b7d
--- /dev/null
+++ b/tests/test_backends.py
@@ -0,0 +1,148 @@
+import pytest
+import numpy as np
+
+from opt_einsum import contract, helpers, contract_expression, backends
+
+try:
+    import tensorflow as tf
+    found_tensorflow = True
+except ImportError:
+    found_tensorflow = False
+
+try:
+    import os
+    os.environ['MKL_THREADING_LAYER'] = 'GNU'
+    import theano
+    found_theano = True
+except ImportError:
+    found_theano = False
+
+try:
+    import cupy
+    found_cupy = True
+except ImportError:
+    found_cupy = False
+
+try:
+    import dask.array as da
+    found_dask = True
+except ImportError:
+    found_dask = False
+
+try:
+    import sparse
+    found_sparse = True
+except ImportError:
+    found_sparse = False
+
+
+tests = [
+    'ab,bc->ca',
+    'abc,bcd,dea',
+    'abc,def->fedcba',
+    # test 'prefer einsum' ops
+    'ijk,ikj',
+    'i,j->ij',
+    'ijk,k->ij',
+]
+
+
+@pytest.mark.skipif(not found_tensorflow, reason="Tensorflow not installed.")
+@pytest.mark.parametrize("string", tests)
+def test_tensorflow(string):
+    views = helpers.build_views(string)
+    ein = contract(string, *views, optimize=False, use_blas=False)
+    opt = np.empty_like(ein)
+
+    shps = [v.shape for v in views]
+    expr = contract_expression(string, *shps, optimize=True)
+
+    sess = tf.Session()
+    with sess.as_default():
+        expr(*views, backend='tensorflow', out=opt)
+
+    assert np.allclose(ein, opt)
+
+    # test non-conversion mode
+    tensorflow_views = backends.convert_arrays_to_tensorflow(views)
+    expr(*tensorflow_views, backend='tensorflow')
+
+
+@pytest.mark.skipif(not found_cupy, reason="Cupy not installed.")
+@pytest.mark.parametrize("string", tests)
+def test_cupy(string):  # pragma: no cover
+    views = helpers.build_views(string)
+    ein = contract(string, *views, optimize=False, use_blas=False)
+    shps = [v.shape for v in views]
+
+    expr = contract_expression(string, *shps, optimize=True)
+
+    opt = expr(*views, backend='cupy')
+    assert np.allclose(ein, opt)
+
+    # test non-conversion mode
+    cupy_views = backends.convert_arrays_to_cupy(views)
+    cupy_opt = expr(*cupy_views, backend='cupy')
+    assert isinstance(cupy_opt, cupy.ndarray)
+    assert np.allclose(ein, cupy.asnumpy(cupy_opt))
+
+
+@pytest.mark.skipif(not found_theano, reason="Theano not installed.")
+@pytest.mark.parametrize("string", tests)
+def test_theano(string):
+    views = helpers.build_views(string)
+    ein = contract(string, *views, optimize=False, use_blas=False)
+    shps = [v.shape for v in views]
+
+    expr = contract_expression(string, *shps, optimize=True)
+
+    opt = expr(*views, backend='theano')
+    assert np.allclose(ein, opt)
+
+    # test non-conversion mode
+    theano_views = backends.convert_arrays_to_theano(views)
+    theano_opt = expr(*theano_views, backend='theano')
+    assert isinstance(theano_opt, theano.tensor.TensorVariable)
+
+
+@pytest.mark.skipif(not found_dask, reason="Dask not installed.")
+@pytest.mark.parametrize("string", tests)
+def test_dask(string):
+    views = helpers.build_views(string)
+    ein = contract(string, *views, optimize=False, use_blas=False)
+    shps = [v.shape for v in views]
+    expr = contract_expression(string, *shps, optimize=True)
+
+    # test non-conversion mode
+    da_views = [da.from_array(x, chunks=(2)) for x in views]
+    da_opt = expr(*da_views, backend='dask')
+
+    # check type is maintained when not using numpy arrays
+    assert isinstance(da_opt, da.Array)
+
+    assert np.allclose(ein, np.array(da_opt))
+
+
+@pytest.mark.skipif(not found_sparse, reason="Sparse not installed.")
+@pytest.mark.parametrize("string", tests)
+def test_sparse(string):
+    views = helpers.build_views(string)
+
+    # sparsify views so they don't become dense during contraction
+    for view in views:
+        np.random.seed(42)
+        mask = np.random.choice([False, True], view.shape, True, [0.05, 0.95])
+        view[mask] = 0
+
+    ein = contract(string, *views, optimize=False, use_blas=False)
+    shps = [v.shape for v in views]
+    expr = contract_expression(string, *shps, optimize=True)
+
+    # test non-conversion mode
+    sparse_views = [sparse.COO.from_numpy(x) for x in views]
+    sparse_opt = expr(*sparse_views, backend='sparse')
+
+    # check type is maintained when not using numpy arrays
+    assert isinstance(sparse_opt, sparse.COO)
+
+    assert np.allclose(ein, sparse_opt.todense())
diff --git a/tests/test_blas.py b/tests/test_blas.py
index 31849b1..1001921 100644
--- a/tests/test_blas.py
+++ b/tests/test_blas.py
@@ -57,11 +57,12 @@ blas_tests = [
    ((['jli', 'ljk'], 'ik', set('lj')),     'TDOT'), # ST GEMM T N Tensor
 
    # Other
-   ((['ijk', 'ikj'], '', set('ijk')),       False), # Transpose DOT
-   ((['ijj', 'jk'], 'ik', set('j')),        False), # Double index
-   ((['i', 'j'], 'ij', set()),              False), # Outer
-   ((['ijk', 'j'], 'ij', set()),            False), # Index sum 1
-   ((['ijk', 'k'], 'ij', set()),            False), # Index sum 2
+   ((['ijk', 'ikj'], '', set('ijk')),       'DOT/EINSUM'  ),  # Transpose DOT
+   ((['i', 'j'], 'ij', set()),              'OUTER/EINSUM'),  # Outer
+   ((['ijk', 'ik'], 'j', set('ik')),        'GEMV/EINSUM' ),  # Matrix-vector
+   ((['ijj', 'jk'], 'ik', set('j')),        False         ),  # Double index
+   ((['ijk', 'j'], 'ij', set()),            False         ),  # Index sum 1
+   ((['ij', 'ij'], 'ij', set()),            False         ),  # Index sum 2
 ]
 
 @pytest.mark.parametrize("inp,benchmark", blas_tests)
diff --git a/tests/test_contract.py b/tests/test_contract.py
index c63f798..f79b910 100644
--- a/tests/test_contract.py
+++ b/tests/test_contract.py
@@ -241,4 +241,4 @@ def test_contract_expression_checks():
     # should only be able to specify out
     with pytest.raises(ValueError) as err:
         expr(np.random.rand(2, 3), np.random.rand(3, 4), order='F')
-    assert "only valid keyword argument to a `ContractExpression`" in str(err)
+    assert "only valid keyword arguments to a `ContractExpression`" in str(err)