Repository: dgasmith/opt_einsum
Pull request #32 (2018-08-04 23:08:54)

add pytorch, tensorflow-eager support, minor optimizations
## Description

There are 3 related changes here:

* The main change is the addition of backend support for [pytorch](https://pytorch.org) and [tensorflow's eager execution mode](https://www.tensorflow.org/guide/eager). Pytorch doesn't have tensordot yet (https://github.com/pytorch/pytorch/issues/8988, https://github.com/pytorch/pytorch/pull/10025), and it's ``transpose`` and ``einsum`` functions also don't match numpy's api so there's wrappers for each of those. All the backends are also in their own submodules now - it should be pretty clear and simple to add any others.

* If two contracted dimensions are *both* size 1 this is now is not marked as a 'broadcast' index, meaning that einsum is not required (good for performance, and also if a backend doesn't supply it).

* A variety of small optimisations, mostly changing for loops w/ list append to comprehensions, and also using ``np.broadcast_to(np.nan, shape)`` for dummy arrays.

-----

Some general notes on my experience of the performance of all the different GPU backends (compilation aside, all can offer pretty drastic speedup over numpy):

**Compiling backends [``tensorflow``, ``theano``]:**

- Compilation is slow (theano even more so than tensorflow) for large contractions, we are talking on the order of seconds
- Once compiled, the expression can be slightly faster to run (maybe 10-15%) than the best of the eager backends below
- Because of these two things - the eager backends below are generally much more preferable unless an expression is going to be extensively reused

**Eager backends [``cupy``, ``pytorch``, ``tensorflow`` in eager mode]:**

- ``cupy`` and ``pytorch`` both seem to have similar performance with eager-``tensorflow`` lagging behind a bit (could very well be specific to the GPU/setup I have)
- ``pytorch`` sadly doesn't have complex dtype support [yet](https://github.com/pytorch/pytorch/issues/755) so is less versatile for quantum/physics applications
- ``tensorflow`` eager mode needs to be enabled before anything else takes place (with ``import tensorflow as tf; tf.enable_eager_execution()``) which is a slight inconvenience

Will all that taken into account, ``cupy`` currently seems to be the easiest, most versatile GPU backend with excellent performance.

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


dgasmith: Looks like torch is getting a segfault in travis. The libraries that we pull down are getting heavy enough that we should consider `conda`. If you do not have objections I can make a PR that switches to using `conda` over `pip`.
jcmgray: No objections at all, ``conda`` certainly makes dependency management a lot easier. In a couple of other projects I install minconda to ``$HOME/conda`` which travis can then cache, speeds up installation a lot.

Patch:
diff --git a/.travis.yml b/.travis.yml
index 3b367bd..af7e375 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,7 +8,6 @@ sudo: false
 matrix:
   include:
     - python: 2.7
-    - python: 3.4
     - python: 3.5
     - python: 3.6
     - python: 3.7
diff --git a/docs/source/backends.rst b/docs/source/backends.rst
index 9f05977..3198ad6 100644
--- a/docs/source/backends.rst
+++ b/docs/source/backends.rst
@@ -2,7 +2,8 @@
 Backends & GPU Support
 ======================
 
-A brief overview of libraries that have been tested:
+The following is a brief overview of libraries which have been tested with
+``opt_einsum``:
 
     - `tensorflow <https://www.tensorflow.org/>`_: compiled tensor expressions
       that can run on GPU.
@@ -13,6 +14,7 @@ A brief overview of libraries that have been tested:
       computations, distributed scheduling, and potential reuse of
       intermediaries.
     - `sparse <https://sparse.pydata.org/>`_: sparse tensors.
+    - `pytorch <https://pytorch.org>`_: numpy-like api for GPU 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
@@ -25,14 +27,14 @@ While more special functionality such as axes reduction is reliant on a
 
 .. note::
 
-    For a contraction to be possible without using a backend einsum, it must 
-    satisfy the following rule: in the full expression (so *including* output 
-    indices) each index must appear twice. In other words, each dimension 
+    For a contraction to be possible without using a backend einsum, it must
+    satisfy the following rule: in the full expression (so *including* output
+    indices) each index must appear twice. In other words, each dimension
     must be contracted with one other dimension, or left alone.
 
 
 General backend for any ndarray
--------------------------------
+===============================
 
 This 'duck-typing' support just requires specifying the correct ``backend``
 argument for the type of arrays supplied when calling
@@ -106,7 +108,7 @@ supported. An example:
 
 
 Special (GPU) backends for numpy arrays
----------------------------------------
+=======================================
 
 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.
@@ -117,12 +119,16 @@ automatically for:
     - `tensorflow <https://www.tensorflow.org/>`_
     - `theano <http://deeplearning.net/software/theano/>`_
     - `cupy <https://cupy.chainer.org/>`_
+    - `pytorch <https://pytorch.org>`_
 
 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.
+Additionally, if arrays are marked as ``constant``
+(see :ref:`constants-section`), then these arrays will be kept on the device
+for optimal performance.
 
 
 Theano
@@ -176,3 +182,46 @@ session:
 Note that you can 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.
+
+Version 1.9 of tensorflow also added support for eager execution of
+computations. If compilation of the contraction expression tensorflow graph is
+taking a substantial amount of time up then it can be advantageous to use this,
+especially since tensor contractions are quite compute-bound. This is achieved
+by running the following snippet:
+
+
+.. code-block:: python
+
+  import tensorflow as tf
+  tf.enable_eager_execution()
+
+After which ``opt_einsum`` will automatically detect eager mode if
+``backend='tensorflow'`` is supplied to a
+:class:`~opt_einsum.ContractExpression`.
+
+
+Pytorch & Cupy
+--------------
+
+Both `pytorch <https://pytorch.org>`_ and `cupy <https://cupy.chainer.org/>`_
+offer numpy-like, GPU-enabled arrays which execute eagerly rather than
+requiring any compilation. If they are installed, no steps are required to
+utilize them other than specifiying the ``backend`` keyword:
+
+.. code-block:: python
+
+    >>> expr(*xs, backend='torch')
+    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)
+
+    >>> expr(*xs, backend='cupy')
+    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)
+
+And as with the other GPU backends, if raw ``cupy`` or ``pytorch`` arrays are
+supplied the returned array will be of the same type, with no conversion
+to or from ``numpy`` arrays.
diff --git a/docs/source/reusing_paths.rst b/docs/source/reusing_paths.rst
index d920d4a..93edc1e 100644
--- a/docs/source/reusing_paths.rst
+++ b/docs/source/reusing_paths.rst
@@ -28,6 +28,8 @@ The ``ContractExpression`` can be called with 3 arrays that match the original s
 Note that few checks are performed when calling the expression, and while it will work for a set of arrays with the same ranks as the original shapes but differing sizes, it might no longer be optimal.
 
 
+.. _constants-section:
+
 ====================
 Specifying Constants
 ====================
@@ -94,9 +96,9 @@ two contractions to compute the output.
 
 .. note::
 
-    The constant part of an expression is lazily generated upon first call,
-    (specific a particular backend) though it can be explicitly built with call
-    to :meth:`~opt_einsum.contract.ContractExpression.evaluate_constants`.
+    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`.
 
 We can confirm the advantage of using expressions and constants by timing the
 following scenarios, first setting
diff --git a/opt_einsum/backends.py b/opt_einsum/backends.py
deleted file mode 100644
index e133d6a..0000000
--- a/opt_einsum/backends.py
+++ /dev/null
@@ -1,216 +0,0 @@
-"""
-Performing array conversions and operations with various backends.
-"""
-
-import importlib
-
-import numpy
-
-# 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 to_tensorflow(array, constant=False):
-    """Convert a numpy array to a ``tensorflow.placeholder`` instance.
-    """
-    import tensorflow
-
-    if isinstance(array, numpy.ndarray):
-        if constant:
-            return tensorflow.constant(array, array.dtype, array.shape)
-
-        return tensorflow.placeholder(array.dtype, array.shape)
-
-    return array
-
-
-def build_tensorflow_expression(arrays, expr):
-    """Build a tensorflow function based on ``arrays`` and ``expr``.
-    """
-    import tensorflow
-
-    placeholders = [to_tensorflow(array) for array in arrays]
-    graph = expr._contract(placeholders, backend='tensorflow')
-
-    def tensorflow_contract(*arrays):
-        session = tensorflow.get_default_session()
-        # only want to feed placeholders - constant tensors already have values
-        feed_dict = {p: a for p, a in zip(placeholders, arrays) if p.op.type == 'Placeholder'}
-        return session.run(graph, feed_dict=feed_dict)
-
-    return tensorflow_contract
-
-
-def evaluate_constants_tensorflow(const_arrays, expr):
-    """Convert constant arguments to tensorflow constants, and perform any
-    possible constant contractions. Requires evaluating a tensorflow graph.
-    """
-    import tensorflow
-
-    # compute the partial graph of new inputs
-    const_arrays = [to_tensorflow(x, constant=True) for x in const_arrays]
-    new_ops, new_contraction_list = expr(*const_arrays, backend='tensorflow', evaluate_constants=True)
-
-    # evaluate the new inputs and convert to tensorflow constants
-    session = tensorflow.get_default_session()
-    new_ops = [None if x is None else to_tensorflow(session.run(x), constant=True) for x in new_ops]
-
-    return new_ops, new_contraction_list
-
-
-# Theano
-
-def to_theano(array, constant=False):
-    """Convert a numpy array to ``theano.tensor.TensorType`` instance.
-    """
-    import theano
-
-    if isinstance(array, numpy.ndarray):
-        if constant:
-            return theano.tensor.constant(array)
-
-        return theano.tensor.TensorType(dtype=array.dtype, broadcastable=[False] * len(array.shape))()
-
-    return array
-
-
-def build_theano_expression(arrays, expr):
-    """Build a theano function based on ``arrays`` and ``expr``.
-    """
-    import theano
-
-    in_vars = [to_theano(array) for array in arrays]
-    out_var = expr._contract(in_vars, backend='theano')
-
-    # don't supply constants to graph
-    graph_ins = [x for x in in_vars if not isinstance(x, theano.tensor.TensorConstant)]
-    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 theano_contract
-
-
-def evaluate_constants_theano(const_arrays, expr):
-    # compute the partial graph of new inputs
-    const_arrays = [to_theano(x, constant=True) for x in const_arrays]
-    new_ops, new_contraction_list = expr(*const_arrays, backend='theano', evaluate_constants=True)
-
-    # evaluate the new inputs and convert to theano shared tensors
-    new_ops = [None if x is None else to_theano(x.eval(), constant=True) for x in new_ops]
-
-    return new_ops, new_contraction_list
-
-
-# Cupy
-
-def to_cupy(array):  # pragma: no cover
-    import cupy
-
-    if isinstance(array, numpy.ndarray):
-        return cupy.asarray(array)
-
-    return array
-
-
-def build_cupy_expression(_, expr):  # pragma: no cover
-    """Build a cupy function based on ``arrays`` and ``expr``.
-    """
-
-    def cupy_contract(*arrays):
-        cupy_arrays = [to_cupy(x) for x in arrays]
-        cupy_out = expr._contract(cupy_arrays, backend='cupy')
-        return cupy_out.get()
-
-    return cupy_contract
-
-
-def evaluate_constants_cupy(const_arrays, expr):  # pragma: no cover
-    """Convert constant arguments to cupy arrays, and perform any possible
-    constant contractions.
-    """
-    const_arrays = [to_cupy(x) for x in const_arrays]
-    return expr(*const_arrays, backend='cupy', evaluate_constants=True)
-
-
-# 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,
-}
-
-
-PARSE_CONSTS_BACKENDS = {
-    'tensorflow': evaluate_constants_tensorflow,
-    'theano': evaluate_constants_theano,
-    'cupy': evaluate_constants_cupy,
-}
-
-
-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)
-
-
-def evaluate_constants(backend, arrays, expr):
-    """Convert constant arrays to the correct backend, and perform as much of
-    the contraction of ``expr`` with these as possible.
-    """
-    return PARSE_CONSTS_BACKENDS[backend](arrays, expr)
diff --git a/opt_einsum/backends/__init__.py b/opt_einsum/backends/__init__.py
new file mode 100644
index 0000000..33e5362
--- /dev/null
+++ b/opt_einsum/backends/__init__.py
@@ -0,0 +1,17 @@
+from .dispatch import (get_func, has_einsum, build_expression, evaluate_constants, CONVERT_BACKENDS)
+from .tensorflow import to_tensorflow
+from .theano import to_theano
+from .cupy import to_cupy
+from .torch import to_torch
+
+__all__ = [
+    'get_func',
+    'has_einsum',
+    'build_expression',
+    'evaluate_constants',
+    'CONVERT_BACKENDS',
+    'to_tensorflow',
+    'to_theano',
+    'to_cupy',
+    'to_torch',
+]
diff --git a/opt_einsum/backends/cupy.py b/opt_einsum/backends/cupy.py
new file mode 100644
index 0000000..1fa4a4f
--- /dev/null
+++ b/opt_einsum/backends/cupy.py
@@ -0,0 +1,32 @@
+"""
+Required functions for optimized contractions of numpy arrays using cupy.
+"""
+
+from __future__ import absolute_import
+import numpy as np
+
+
+def to_cupy(array):  # pragma: no cover
+    import cupy
+
+    if isinstance(array, np.ndarray):
+        return cupy.asarray(array)
+
+    return array
+
+
+def build_expression(_, expr):  # pragma: no cover
+    """Build a cupy function based on ``arrays`` and ``expr``.
+    """
+
+    def cupy_contract(*arrays):
+        return expr._contract([to_cupy(x) for x in arrays], backend='cupy').get()
+
+    return cupy_contract
+
+
+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)
diff --git a/opt_einsum/backends/dispatch.py b/opt_einsum/backends/dispatch.py
new file mode 100644
index 0000000..f766f3f
--- /dev/null
+++ b/opt_einsum/backends/dispatch.py
@@ -0,0 +1,103 @@
+"""
+Handles dispatching array operations to the correct backend library, as well
+as converting arrays to backend formats and then potentially storing them as
+constants.
+"""
+
+import importlib
+
+import numpy
+
+from . import cupy as _cupy
+from . import torch as _torch
+from . import theano as _theano
+from . import tensorflow as _tensorflow
+
+# known non top-level imports
+_aliases = {
+    'dask': 'dask.array',
+    'theano': 'theano.tensor',
+    'torch': 'opt_einsum.backends.torch',
+}
+
+
+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]
+
+
+# Dispatch to correct expression backend
+#    these are the backends which support explicit to-and-from numpy conversion
+CONVERT_BACKENDS = {
+    'tensorflow': _tensorflow.build_expression,
+    'theano': _theano.build_expression,
+    'cupy': _cupy.build_expression,
+    'torch': _torch.build_expression,
+}
+
+
+EVAL_CONSTS_BACKENDS = {
+    'tensorflow': _tensorflow.evaluate_constants,
+    'theano': _theano.evaluate_constants,
+    'cupy': _cupy.evaluate_constants,
+    'torch': _torch.evaluate_constants,
+}
+
+
+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)
+
+
+def evaluate_constants(backend, arrays, expr):
+    """Convert constant arrays to the correct backend, and perform as much of
+    the contraction of ``expr`` with these as possible.
+    """
+    return EVAL_CONSTS_BACKENDS[backend](arrays, expr)
diff --git a/opt_einsum/backends/tensorflow.py b/opt_einsum/backends/tensorflow.py
new file mode 100644
index 0000000..a6d218b
--- /dev/null
+++ b/opt_einsum/backends/tensorflow.py
@@ -0,0 +1,119 @@
+"""
+Required functions for optimized contractions of numpy arrays using tensorflow.
+"""
+
+from __future__ import absolute_import
+import numpy as np
+
+
+_CACHED_TF_DEVICE = None
+
+
+def _get_tensorflow_and_device():
+    global _CACHED_TF_DEVICE
+
+    if _CACHED_TF_DEVICE is None:
+        import tensorflow as tf
+
+        try:
+            eager = tf.contrib.eager.in_eager_mode()
+        except AttributeError:
+            eager = False
+
+        device = tf.test.gpu_device_name()
+        if not device:
+            device = 'cpu'
+
+        _CACHED_TF_DEVICE = tf, device, eager
+
+    return _CACHED_TF_DEVICE
+
+
+def to_tensorflow(array, constant=False):
+    """Convert a numpy array to a ``tensorflow.placeholder`` instance.
+    """
+    tf, device, eager = _get_tensorflow_and_device()
+
+    if eager:
+        if isinstance(array, np.ndarray):
+            with tf.device(device):
+                return tf.convert_to_tensor(array)
+
+        return array
+
+    if isinstance(array, np.ndarray):
+        if constant:
+            return tf.convert_to_tensor(array)
+
+        return tf.placeholder(array.dtype, array.shape)
+
+    return array
+
+
+# Standard graph mode
+
+def build_expression_graph(arrays, expr):
+    """Build a tensorflow function based on ``arrays`` and ``expr``.
+    """
+    tf, _, _ = _get_tensorflow_and_device()
+
+    placeholders = [to_tensorflow(array) for array in arrays]
+    graph = expr._contract(placeholders, backend='tensorflow')
+
+    def tensorflow_contract(*arrays):
+        session = tf.get_default_session()
+        # only want to feed placeholders - constant tensors already have values
+        feed_dict = {p: a for p, a in zip(placeholders, arrays) if p.op.type == 'Placeholder'}
+        return session.run(graph, feed_dict=feed_dict)
+
+    return tensorflow_contract
+
+
+def evaluate_constants_graph(const_arrays, expr):
+    """Convert constant arguments to tensorflow constants, and perform any
+    possible constant contractions. Requires evaluating a tensorflow graph.
+    """
+    tf, _, _ = _get_tensorflow_and_device()
+
+    # compute the partial graph of new inputs
+    const_arrays = [to_tensorflow(x, constant=True) for x in const_arrays]
+    new_ops, new_contraction_list = expr(*const_arrays, backend='tensorflow', evaluate_constants=True)
+
+    # evaluate the new inputs and convert back to tensorflow, maintaining None as non-consts
+    session = tf.get_default_session()
+    new_consts = iter(session.run([x for x in new_ops if x is not None]))
+    new_ops = [None if x is None else to_tensorflow(next(new_consts), constant=True) for x in new_ops]
+
+    return new_ops, new_contraction_list
+
+
+# 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()
+
+    return tensorflow_eager_contract
+
+
+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)
+
+
+# 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
+    return fn(arrays, expr)
+
+
+def evaluate_constants(const_arrays, expr):
+    _, _, eager = _get_tensorflow_and_device()
+    fn = evaluate_constants_eager if eager else evaluate_constants_graph
+    return fn(const_arrays, expr)
diff --git a/opt_einsum/backends/theano.py b/opt_einsum/backends/theano.py
new file mode 100644
index 0000000..80c0bef
--- /dev/null
+++ b/opt_einsum/backends/theano.py
@@ -0,0 +1,49 @@
+"""
+Required functions for optimized contractions of numpy arrays using theano.
+"""
+
+from __future__ import absolute_import
+import numpy as np
+
+
+def to_theano(array, constant=False):
+    """Convert a numpy array to ``theano.tensor.TensorType`` instance.
+    """
+    import theano
+
+    if isinstance(array, np.ndarray):
+        if constant:
+            return theano.tensor.constant(array)
+
+        return theano.tensor.TensorType(dtype=array.dtype, broadcastable=[False] * len(array.shape))()
+
+    return array
+
+
+def build_expression(arrays, expr):
+    """Build a theano function based on ``arrays`` and ``expr``.
+    """
+    import theano
+
+    in_vars = [to_theano(array) for array in arrays]
+    out_var = expr._contract(in_vars, backend='theano')
+
+    # don't supply constants to graph
+    graph_ins = [x for x in in_vars if not isinstance(x, theano.tensor.TensorConstant)]
+    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 theano_contract
+
+
+def evaluate_constants(const_arrays, expr):
+    # compute the partial graph of new inputs
+    const_arrays = [to_theano(x, constant=True) for x in const_arrays]
+    new_ops, new_contraction_list = expr(*const_arrays, backend='theano', evaluate_constants=True)
+
+    # evaluate the new inputs and convert to theano shared tensors
+    new_ops = [None if x is None else to_theano(x.eval(), constant=True) for x in new_ops]
+
+    return new_ops, new_contraction_list
diff --git a/opt_einsum/backends/torch.py b/opt_einsum/backends/torch.py
new file mode 100644
index 0000000..698a4a5
--- /dev/null
+++ b/opt_einsum/backends/torch.py
@@ -0,0 +1,116 @@
+"""
+Required functions for optimized contractions of numpy arrays using pytorch.
+"""
+
+from __future__ import absolute_import
+import numpy as np
+
+from ..parser import einsum_symbols_base
+
+
+_TORCH_DEVICE = None
+
+
+def _get_torch_and_device():
+    global _TORCH_DEVICE
+
+    if _TORCH_DEVICE is None:
+        import torch
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        _TORCH_DEVICE = torch, device
+
+    return _TORCH_DEVICE
+
+
+def transpose(a, axes):
+    """Normal torch transpose is only valid for 2D matrices.
+    """
+    return a.permute(*axes)
+
+
+def einsum(equation, *operands):
+    """Variadic version of torch.einsum to match numpy api.
+    """
+    torch, _ = _get_torch_and_device()
+    return torch.einsum(equation, operands)
+
+
+def tensordot(x, y, axes=2):
+    """Simple translation of tensordot syntax to einsum.
+    """
+    # XXX: tensordot should be directly implemented in torch soon
+    torch, _ = _get_torch_and_device()
+
+    xnd = x.ndimension()
+    ynd = y.ndimension()
+
+    # convert int argument to (list[int], list[int])
+    if isinstance(axes, int):
+        axes = range(xnd - axes, xnd), range(axes)
+
+    # convert (int, int) to (list[int], list[int])
+    if isinstance(axes[0], int):
+        axes = (axes[0],), axes[1]
+    if isinstance(axes[1], int):
+        axes = axes[0], (axes[1],)
+
+    # initialize empty indices
+    x_ix = [None] * xnd
+    y_ix = [None] * ynd
+    out_ix = []
+
+    # fill in repeated indices
+    available_ix = iter(einsum_symbols_base)
+    for ax1, ax2 in zip(*axes):
+        repeat = next(available_ix)
+        x_ix[ax1] = repeat
+        y_ix[ax2] = repeat
+
+    # fill in the rest, and maintain output order
+    for i in range(xnd):
+        if x_ix[i] is None:
+            leave = next(available_ix)
+            x_ix[i] = leave
+            out_ix.append(leave)
+    for i in range(ynd):
+        if y_ix[i] is None:
+            leave = next(available_ix)
+            y_ix[i] = leave
+            out_ix.append(leave)
+
+    # form full string and contract!
+    einsum_str = "{},{}->{}".format(*map("".join, (x_ix, y_ix, out_ix)))
+    return torch.einsum(einsum_str, (x, y))
+
+
+def to_torch(array):
+    torch, device = _get_torch_and_device()
+
+    if isinstance(array, np.ndarray):
+        return torch.from_numpy(array).to(device)
+
+    return array
+
+
+def build_expression(_, expr):  # pragma: no cover
+    """Build a torch function based on ``arrays`` and ``expr``.
+    """
+
+    def torch_contract(*arrays):
+        torch_arrays = [to_torch(x) for x in arrays]
+        torch_out = expr._contract(torch_arrays, backend='torch')
+
+        if torch_out.device.type == 'cpu':
+            return torch_out.numpy()
+
+        return torch_out.cpu().numpy()
+
+    return torch_contract
+
+
+def evaluate_constants(const_arrays, expr):
+    """Convert constant arguments to torch, and perform any possible constant
+    contractions.
+    """
+    const_arrays = [to_torch(x) for x in const_arrays]
+    return expr(*const_arrays, backend='torch', evaluate_constants=True)
diff --git a/opt_einsum/blas.py b/opt_einsum/blas.py
index b8170ce..5ace3b6 100644
--- a/opt_einsum/blas.py
+++ b/opt_einsum/blas.py
@@ -7,7 +7,7 @@ import numpy as np
 from . import helpers
 
 
-def can_blas(inputs, result, idx_removed):
+def can_blas(inputs, result, idx_removed, shapes=None):
     """
     Checks if we can use a BLAS call.
 
@@ -19,7 +19,8 @@ def can_blas(inputs, result, idx_removed):
         Resulting summation.
     idx_removed : set
         Indices that are removed in the summation
-
+    shapes : sequence of tuple[int], optional
+        If given, check also that none of the indices are broadcast dimensions.
 
     Returns
     -------
@@ -35,14 +36,18 @@ def can_blas(inputs, result, idx_removed):
 
     Examples
     --------
-    >>> _can_blas(['ij', 'jk'], 'ik', set('j'))
+    >>> can_blas(['ij', 'jk'], 'ik', set('j'))
     'GEMM'
 
-    >>> _can_blas(['ijj', 'jk'], 'ik', set('j'))
+    >>> can_blas(['ijj', 'jk'], 'ik', set('j'))
     False
 
-    >>> _can_blas(['ab', 'cd'], 'abcd', set())
+    >>> can_blas(['ab', 'cd'], 'abcd', set())
     'OUTER/EINSUM'
+
+    >>> # looks like GEMM but actually 'j' is broadcast:
+    >>> can_blas(['ij', 'jk'], 'ik', set('j'), shapes=[(4, 1), (5, 6)])
+    False
     """
     # Can only do two
     if len(inputs) != 2:
@@ -62,6 +67,13 @@ def can_blas(inputs, result, idx_removed):
         if nl + nr - 1 == int(c in result):
             return False
 
+    # check for broadcast indices e.g:
+    #     "ij,jk->ik" (but one of the 'j' dimensions is broadcast up)
+    if shapes is not None:
+        for c in idx_removed:
+            if shapes[0][input_left.find(c)] != shapes[1][input_right.find(c)]:
+                return False
+
     # Prefer einsum if not removing indices
     #     (N.B. tensordot outer faster for large arrays?)
     if len(idx_removed) == 0:
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index e1fc026..f387699 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -123,14 +123,14 @@ def contract_path(*operands, **kwargs):
     # Build a few useful list and sets
     input_list = input_subscripts.split(',')
     input_sets = [set(x) for x in input_list]
+    input_shps = [x.shape for x in operands]
     output_set = set(output_subscript)
     indices = set(input_subscripts.replace(',', ''))
 
     # Get length of each unique dimension and ensure all dimensions are correct
     dimension_dict = {}
-    broadcast_indices = [[] for x in range(len(input_list))]
     for tnum, term in enumerate(input_list):
-        sh = operands[tnum].shape
+        sh = input_shps[tnum]
 
         if len(sh) != len(term):
             raise ValueError("Einstein sum subscript %s does not contain the "
@@ -138,11 +138,7 @@ def contract_path(*operands, **kwargs):
         for cnum, char in enumerate(term):
             dim = sh[cnum]
 
-            # Build out broadcast indices
-            if dim == 1:
-                broadcast_indices[tnum].append(char)
-
-            if char in dimension_dict.keys():
+            if char in dimension_dict:
                 # For broadcasting cases we always want the largest dim size
                 if dimension_dict[char] == 1:
                     dimension_dict[char] = dim
@@ -152,13 +148,8 @@ def contract_path(*operands, **kwargs):
             else:
                 dimension_dict[char] = dim
 
-    # Convert broadcast inds to sets
-    broadcast_indices = [set(x) for x in broadcast_indices]
-
     # Compute size of each input array plus the output array
-    size_list = []
-    for term in input_list + [output_subscript]:
-        size_list.append(helpers.compute_size_by_dict(term, dimension_dict))
+    size_list = [helpers.compute_size_by_dict(term, dimension_dict) for term in input_list + [output_subscript]]
     out_size = max(size_list)
 
     if memory_limit is None:
@@ -217,17 +208,11 @@ def contract_path(*operands, **kwargs):
         scale_list.append(len(idx_contract))
         size_list.append(helpers.compute_size_by_dict(out_inds, dimension_dict))
 
-        bcast = set()
-        tmp_inputs = []
-        for x in contract_inds:
-            tmp_inputs.append(input_list.pop(x))
-            bcast |= broadcast_indices.pop(x)
-
-        new_bcast_inds = bcast - idx_removed
+        tmp_inputs = [input_list.pop(x) for x in contract_inds]
+        tmp_shapes = [input_shps.pop(x) for x in contract_inds]
 
-        # If were broadcasting, nix blas
-        if use_blas and not len(idx_removed & bcast):
-            do_blas = blas.can_blas(tmp_inputs, out_inds, idx_removed)
+        if use_blas:
+            do_blas = blas.can_blas(tmp_inputs, out_inds, idx_removed, tmp_shapes)
         else:
             do_blas = False
 
@@ -235,11 +220,15 @@ def contract_path(*operands, **kwargs):
         if (cnum - len(path)) == -1:
             idx_result = output_subscript
         else:
-            sort_result = [(dimension_dict[ind], ind) for ind in out_inds]
-            idx_result = "".join([x[1] for x in sorted(sort_result)])
+            # use tensordot order to minimize transpositions
+            all_input_inds = "".join(tmp_inputs)
+            idx_result = "".join(sorted(out_inds, key=all_input_inds.find))
+
+        shp_result = parser.find_output_shape(tmp_inputs, tmp_shapes, idx_result)
 
         input_list.append(idx_result)
-        broadcast_indices.append(new_bcast_inds)
+        input_shps.append(shp_result)
+
         einsum_str = ",".join(tmp_inputs) + "->" + idx_result
 
         contraction = (contract_inds, idx_removed, einsum_str, input_list[:], do_blas)
@@ -301,7 +290,7 @@ def _transpose(x, axes, backend='numpy'):
     """
     try:
         return x.transpose(axes)
-    except AttributeError:
+    except (AttributeError, TypeError):
         # some libraries don't implement method version
         fn = backends.get_func('transpose', backend)
         return fn(x, axes)
@@ -674,13 +663,11 @@ class ContractExpression:
         return s
 
 
-class _ShapeOnly(np.ndarray):
+def shape_only(shape):
     """Dummy ``numpy.ndarray`` which has a shape only - for generating
     contract expressions.
     """
-
-    def __init__(self, shape):
-        self.shape = shape
+    return np.broadcast_to(np.nan, shape)
 
 
 def contract_expression(subscripts, *shapes, **kwargs):
@@ -760,6 +747,6 @@ def contract_expression(subscripts, *shapes, **kwargs):
     kwargs['_constants_dict'] = constants_dict
 
     # apart from constant arguments, make dummy arrays
-    dummy_arrays = [s if i in constants else _ShapeOnly(s) for i, s in enumerate(shapes)]
+    dummy_arrays = [s if i in constants else shape_only(s) for i, s in enumerate(shapes)]
 
     return contract(subscripts, *dummy_arrays, **kwargs)
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index 865706e..786e5ae 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -76,11 +76,17 @@ def find_output_str(subscripts):
     """Find the output string for the inputs ``susbcripts``.
     """
     tmp_subscripts = subscripts.replace(",", "")
-    output_subscript = ""
-    for s in sorted(set(tmp_subscripts)):
-        if tmp_subscripts.count(s) == 1:
-            output_subscript += s
-    return output_subscript
+    return "".join(s for s in sorted(set(tmp_subscripts)) if tmp_subscripts.count(s) == 1)
+
+
+def find_output_shape(inputs, shapes, output):
+    """Find the output shape for given inputs, shapes and output string, taking
+    into account broadcasting.
+    """
+    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
+    )
 
 
 def possibly_convert_to_numpy(x):
@@ -111,10 +117,10 @@ def parse_einsum_input(operands):
 
     >>> a = np.random.rand(4, 4)
     >>> b = np.random.rand(4, 4, 4)
-    >>> _parse_einsum_input(('...a,...a->...', a, b))
+    >>> parse_einsum_input(('...a,...a->...', a, b))
     ('za,xza', 'xz', [a, b])
 
-    >>> _parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0]))
+    >>> parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0]))
     ('za,xza', 'xz', [a, b])
     """
 
@@ -144,7 +150,7 @@ def parse_einsum_input(operands):
                 elif isinstance(s, int):
                     subscripts += get_symbol(s)
                 else:
-                    raise TypeError("For this input type lists must contain " "either int or Ellipsis")
+                    raise TypeError("For this input type lists must contain either int or Ellipsis")
             if num != last:
                 subscripts += ","
 
@@ -156,7 +162,8 @@ def parse_einsum_input(operands):
                 elif isinstance(s, int):
                     subscripts += get_symbol(s)
                 else:
-                    raise TypeError("For this input type lists must contain " "either int or Ellipsis")
+                    raise TypeError("For this input type lists must contain either int or Ellipsis")
+
     # Check for proper "->"
     if ("-" in subscripts) or (">" in subscripts):
         invalid = (subscripts.count("-") > 1) or (subscripts.count(">") > 1)


Test patch:
diff --git a/devtools/conda-envs/full_test.yaml b/devtools/conda-envs/full_test.yaml
index d9ad304..412d3cb 100644
--- a/devtools/conda-envs/full_test.yaml
+++ b/devtools/conda-envs/full_test.yaml
@@ -1,6 +1,7 @@
 name: test
 channels:
   - conda-forge
+  - pytorch
 dependencies:
     # Base depends
   - numpy
@@ -11,6 +12,7 @@ dependencies:
   - dask
   - sparse
   - theano
+  - pytorch-cpu
 
     # Testing
   - pytest
diff --git a/opt_einsum/tests/test_backends.py b/opt_einsum/tests/test_backends.py
index 19d72c2..3f62ad1 100644
--- a/opt_einsum/tests/test_backends.py
+++ b/opt_einsum/tests/test_backends.py
@@ -5,8 +5,10 @@ from opt_einsum import contract, helpers, contract_expression, backends
 
 try:
     import tensorflow as tf
+    # needed so tensorflow doesn't allocate all gpu mem
+    _TF_CONFIG = tf.ConfigProto()
+    _TF_CONFIG.gpu_options.allow_growth = True
     found_tensorflow = True
-    sess = tf.Session()
 except ImportError:
     found_tensorflow = False
 
@@ -19,23 +21,10 @@ except ImportError:
     found_theano = False
 
 try:
-    import cupy
-    found_cupy = True
+    import torch
+    found_torch = 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
-
+    found_torch = False
 
 tests = [
     'ab,bc->ca',
@@ -58,8 +47,10 @@ def test_tensorflow(string):
     shps = [v.shape for v in views]
     expr = contract_expression(string, *shps, optimize=True)
 
+    sess = tf.Session(config=_TF_CONFIG)
     with sess.as_default():
         expr(*views, backend='tensorflow', out=opt)
+    sess.close()
 
     assert np.allclose(ein, opt)
 
@@ -81,8 +72,10 @@ def test_tensorflow_with_constants():
     expr = contract_expression(eq, *ops, constants=constants)
 
     # check tensorflow
+    sess = tf.Session(config=_TF_CONFIG)
     with sess.as_default():
         res_got = expr(var, backend='tensorflow')
+    sess.close()
     assert 'tensorflow' in expr._evaluated_constants
     assert np.allclose(res_exp, res_got)
 
@@ -139,9 +132,10 @@ def test_theano_with_constants():
     assert isinstance(res_got3, theano.tensor.TensorVariable)
 
 
-@pytest.mark.skipif(not found_cupy, reason="Cupy not installed.")
 @pytest.mark.parametrize("string", tests)
 def test_cupy(string):  # pragma: no cover
+    cupy = pytest.importorskip("cupy")
+
     views = helpers.build_views(string)
     ein = contract(string, *views, optimize=False, use_blas=False)
     shps = [v.shape for v in views]
@@ -158,8 +152,9 @@ def test_cupy(string):  # pragma: no cover
     assert np.allclose(ein, cupy.asnumpy(cupy_opt))
 
 
-@pytest.mark.skipif(not found_cupy, reason="Cupy not installed.")
-def test_cupy_with_constants():
+def test_cupy_with_constants():  # pragma: no cover
+    cupy = pytest.importorskip("cupy")
+
     eq = 'ij,jk,kl->li'
     shapes = (2, 3), (3, 4), (4, 5)
     constants = {0, 2}
@@ -185,9 +180,10 @@ def test_cupy_with_constants():
     assert np.allclose(res_exp, res_got3.get())
 
 
-@pytest.mark.skipif(not found_dask, reason="Dask not installed.")
 @pytest.mark.parametrize("string", tests)
 def test_dask(string):
+    da = pytest.importorskip("dask.array")
+
     views = helpers.build_views(string)
     ein = contract(string, *views, optimize=False, use_blas=False)
     shps = [v.shape for v in views]
@@ -208,9 +204,10 @@ def test_dask(string):
     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):
+    sparse = pytest.importorskip("sparse")
+
     views = helpers.build_views(string)
 
     # sparsify views so they don't become dense during contraction
@@ -236,3 +233,52 @@ def test_sparse(string):
     sparse_opt = contract(string, *sparse_views, backend='sparse')
     assert isinstance(sparse_opt, sparse.COO)
     assert np.allclose(ein, sparse_opt.todense())
+
+
+@pytest.mark.skipif(not found_torch, reason="Torch not installed.")
+@pytest.mark.parametrize("string", tests)
+def test_torch(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='torch')
+    assert np.allclose(ein, opt)
+
+    # test non-conversion mode
+    torch_views = [backends.to_torch(view) for view in views]
+    torch_opt = expr(*torch_views, backend='torch')
+    assert isinstance(torch_opt, torch.Tensor)
+    assert np.allclose(ein, torch_opt.cpu().numpy())
+
+
+@pytest.mark.skipif(not found_torch, reason="Torch not installed.")
+def test_torch_with_constants():
+
+    eq = 'ij,jk,kl->li'
+    shapes = (2, 3), (3, 4), (4, 5)
+    constants = {0, 2}
+    ops = [np.random.rand(*shp) if i in constants else shp for i, shp in enumerate(shapes)]
+    var = np.random.rand(*shapes[1])
+
+    res_exp = contract(eq, ops[0], var, ops[2])
+
+    expr = contract_expression(eq, *ops, constants=constants)
+
+    # check torch
+    res_got = expr(var, backend='torch')
+    assert 'torch' in expr._evaluated_constants
+    assert np.allclose(res_exp, res_got)
+
+    # check can call with numpy still
+    res_got2 = expr(var, backend='numpy')
+    assert np.allclose(res_exp, res_got2)
+
+    # check torch call returns torch still
+    res_got3 = expr(backends.to_torch(var), backend='torch')
+    assert isinstance(res_got3, torch.Tensor)
+    res_got3 = res_got3.numpy() if res_got3.device.type == 'cpu' else res_got3.cpu().numpy()
+    assert np.allclose(res_exp, res_got3)
diff --git a/opt_einsum/tests/test_edge_cases.py b/opt_einsum/tests/test_edge_cases.py
index 9ee5824..1e6a3d9 100644
--- a/opt_einsum/tests/test_edge_cases.py
+++ b/opt_einsum/tests/test_edge_cases.py
@@ -2,12 +2,10 @@
 Tets a series of opt_einsum contraction paths to ensure the results are the same for different paths
 """
 
-import sys
-
 import numpy as np
 import pytest
 
-from opt_einsum import contract, contract_path, helpers, contract_expression
+from opt_einsum import contract, contract_expression
 
 
 def test_contract_expression_checks():
@@ -114,3 +112,14 @@ def test_broadcasting_contraction4():
     opt = contract('obk,ijk->ioj', a, a, optimize=True)
 
     assert np.allclose(ein, opt)
+
+
+def test_can_blas_on_healed_broadcast_dimensions():
+
+    expr = contract_expression("ab,bc,bd->acd", (5, 4), (1, 5), (4, 20))
+    # first contraction involves broadcasting
+    assert expr.contraction_list[0][2] == 'bc,ab->bca'
+    assert expr.contraction_list[0][-1] is False
+    # but then is healed GEMM is usable
+    assert expr.contraction_list[1][2] == 'bca,bd->acd'
+    assert expr.contraction_list[1][-1] == 'GEMM'
diff --git a/opt_einsum/tests/test_input.py b/opt_einsum/tests/test_input.py
index 3984142..d44b5dc 100644
--- a/opt_einsum/tests/test_input.py
+++ b/opt_einsum/tests/test_input.py
@@ -220,6 +220,7 @@ def test_singleton_dimension_broadcast():
 
     p = np.ones((1, 5))
     q = np.ones((5, 5))
+
     for optimize in (True, False):
         res1 = contract("...ij,...jk->...ik", p, p, optimize=optimize),
         res2 = contract("...ij,...jk->...ik", p, q, optimize=optimize)