Repository: google/jax

Issue:
np.einsum support
support for generic tensor contractions would cover a large class of computations and also provide a foundation for higher order operations. Perhaps jax could then also be added as a `opt_einsum` backend?


Discussion:
Yes, definitely! And you're right that opt_einsum should be very useful here.

Broadly speaking I think we would want to use opt_einsum to translate an `einsum` call into a sequence of calls into `lax.dot_general`, which models the [XLA DotGeneral HLO](https://www.tensorflow.org/xla/operation_semantics#dotgeneral). But to support all of np.einsum's functionality we'd need to handle a few more cases that aren't exactly tensor contractions, like `np.einsum('ii->i', A)` for extracting a matrix diagonal, `np.einsum('ii->', A)` for computing a trace, and `np.einsum('ij->', A)` for a sum over axes (contracting against an implicit all-ones array, you could say).

EDIT: There are also a few annoying constraints in the current DotGeneral HLO, like [the requirement that batch dimensions precede everything else](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/shape_inference.cc#L664), and [the requirement that there is only one contracting dimension per call](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/shape_inference.cc#L647), which we'd have to work around in our lax_numpy.py implementation. I believe the XLA folks are planning to make DotGeneral more convenient to use in the near future, but I'm not sure when.

Let's leave this issue open to track this feature. We can probably also implement `np.tensordot` support at the same time.

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -16,27 +16,29 @@
 from __future__ import division
 from __future__ import print_function
 
-from six.moves import builtins
+import collections
+import itertools
+import string
 
-import six
 import numpy as onp
+import opt_einsum
+import six
+from six.moves import builtins
 
+from jax import jit
 from .. import core
 from ..abstract_arrays import UnshapedArray, ShapedArray, ConcreteArray
 from ..interpreters.xla import DeviceArray
 from .. import lax
-from ..util import memoize, partial, get_module_functions, prod as _prod
+from ..util import memoize, partial, get_module_functions, unzip2, prod as _prod
 from ..lib import xla_bridge
 
-# To provide the same module-level names as Numpy, we need to redefine builtins
-# and also use some common names (like 'shape' and 'dtype') at the top-level.
-# pylint: disable=redefined-builtin,redefined-outer-name
-
-# There might be a pylint bug with tuple unpacking.
-# pylint: disable=unbalanced-tuple-unpacking
-
-# We get docstrings from the underlying numpy functions.
-# pylint: disable=missing-docstring
+if six.PY3:
+  def removechars(s, chars):
+    return s.translate(str.maketrans(dict.fromkeys(chars)))
+else:
+  def removechars(s, chars):
+    return s.translate(None, ''.join(chars))
 
 
 # We replace some builtin names to follow Numpy's API, so we capture here.
@@ -1054,6 +1056,172 @@ def tensordot(a, b, axes=2):
   raise TypeError(msg)
 
 
+def einsum(*operands):
+  # using einsum_call=True here is an internal api for opt_einsum
+  operands, contractions = opt_einsum.contract_path(
+      *operands, einsum_call=True, use_blas=True)
+  contractions = tuple(data[:3] for data in contractions)
+  return _einsum(operands, contractions)
+
+
+@partial(jit, static_argnums=(1,))
+def _einsum(operands, contractions):
+  operands = list(_promote_dtypes(*operands))
+  sum = lambda x, axes: lax.reduce(x, onp.array(0, x.dtype), lax.add, axes)
+
+  def sum_uniques(operand, names, uniques):
+    if uniques:
+      axes = [names.index(name) for name in uniques]
+      operand = sum(operand, axes)
+      names = removechars(names, uniques)
+    return operand, names
+
+  def sum_repeats(operand, names, counts, keep_names):
+    for name, count in counts.items():
+      if count > 1:
+        axes = [i for i, n in enumerate(names) if n == name]
+        eye = lax.broadcasted_eye(operand.dtype, operand.shape, axes)
+        if name not in keep_names:
+          operand = sum(operand * eye, axes)
+          names = names.replace(name, '')
+        else:
+          operand = sum(operand * eye, axes[:-1])
+          names = names.replace(name, '', count - 1)
+    return operand, names
+
+  for operand_indices, contracted_names, einstr in contractions:
+    input_str, result_names = einstr.split('->')
+    input_names = input_str.split(',')
+
+    # switch on the number of operands to be processed in this loop iteration.
+    # every case here sets 'result' and 'names'.
+    if len(operand_indices) == 1:
+      operand = operands.pop(operand_indices[0])
+      names, = input_names
+      counts = collections.Counter(names)
+
+      # sum out unique contracted indices with a single reduce-sum
+      uniques = [name for name in contracted_names if counts[name] == 1]
+      operand, names = sum_uniques(operand, names, uniques)
+
+      # for every repeated index, do a contraction against an identity matrix
+      operand, names = sum_repeats(operand, names, counts, result_names)
+
+    elif len(operand_indices) == 2:
+      lhs, rhs = map(operands.pop, operand_indices)
+      lhs_counts, rhs_counts = map(collections.Counter, input_names)
+      lhs_names, rhs_names = input_names
+
+      # sum out unique contracted indices in lhs and rhs
+      lhs_uniques = [name for name in contracted_names
+                     if lhs_counts[name] == 1 and rhs_counts[name] == 0]
+      lhs, lhs_names = sum_uniques(lhs, lhs_names, lhs_uniques)
+
+      rhs_uniques = [name for name in contracted_names
+                     if rhs_counts[name] == 1 and lhs_counts[name] == 0]
+      rhs, rhs_names = sum_uniques(rhs, rhs_names, rhs_uniques)
+
+      # for every repeated index, contract against an identity matrix
+      lhs, lhs_names = sum_repeats(lhs, lhs_names, lhs_counts,
+                                   result_names + rhs_names)
+      rhs, rhs_names = sum_repeats(rhs, rhs_names, rhs_counts,
+                                   result_names + lhs_names)
+
+      contracted_names = contracted_names & (set(lhs_names) | set(rhs_names))
+      batch_names = set(lhs_names) & set(rhs_names) - contracted_names
+      lhs_batch, rhs_batch = unzip2((lhs_names.find(n), rhs_names.find(n))
+                                    for n in batch_names)
+      if contracted_names:
+        # contract usint lax.dot_general
+        lhs_cont, rhs_cont = unzip2((lhs_names.index(n), rhs_names.index(n))
+                                    for n in contracted_names)
+
+        # lax.dot_general batch dims have to precede non-batch dims
+        batch_dims = tuple(range(len(batch_names)))
+        if lhs_batch != rhs_batch or set(lhs_batch) != set(batch_dims):
+          lhs = moveaxis(lhs, lhs_batch, batch_dims)
+          rhs = moveaxis(rhs, rhs_batch, batch_dims)
+          batch_names = ''.join(batch_names)
+        else:
+          batch_dims = tuple(lhs_batch)
+          batch_names = ''.join(lhs_names[i] for i in batch_dims)
+
+        operand = _dot_general(lhs, rhs, lhs_cont, rhs_cont, len(batch_dims))
+        deleted_names = batch_names + ''.join(contracted_names)
+        names = (batch_names + removechars(lhs_names, deleted_names)
+                 + removechars(rhs_names, deleted_names))
+      else:
+        # no contraction, just a tensor product
+        if lhs_batch != rhs_batch:
+          rhs = moveaxis(rhs, rhs_batch, lhs_batch)
+        batch_names = ''.join(lhs_names[i] for i in lhs_batch)
+
+        names = batch_names + lhs_names + rhs_names
+        lhs_shape = iter(lhs.shape)
+        lhs_shape = [next(lhs_shape) if n in batch_names + lhs_names else 1
+                     for n in names]
+        rhs_shape = iter(rhs.shape)
+        rhs_shape = [next(rhs_shape) if n in batch_names + rhs_names else 1
+                     for n in names]
+        operand = lax.reshape(lhs, lhs_shape) * lax.reshape(rhs, rhs_shape)
+
+    else:
+      raise NotImplementedError
+
+    # the resulting 'operand' with axis labels 'names' should be a permutation
+    # of the desired result
+    assert len(names) == len(result_names) == len(set(names))
+    assert set(names) == set(result_names)
+    if names != result_names:
+      perm = tuple([names.index(name) for name in result_names])
+      operand = lax.transpose(operand, perm)
+    operands.append(operand)  # used in next iteration
+
+  return operands[0]
+
+
+def _dot_general(lhs, rhs, lhs_cont, rhs_cont, nbatch):
+  """Helper for einsum contractions."""
+  # lax.dot_general has some tight constraints on dimension_numbers that this
+  # wrapper loosens via transposes and reshapes
+  assert len(lhs_cont) == len(rhs_cont) > 0
+  ncont = len(lhs_cont)
+  lhs_ntensor = lhs.ndim - nbatch - ncont
+  rhs_ntensor = rhs.ndim - nbatch - ncont
+  batch_dims = tuple(range(nbatch))
+
+  if ncont == 1 and 0 <= lhs_ntensor <= 1 and 0 <= rhs_ntensor <= 1:
+    dimension_numbers = [(lhs_cont, rhs_cont), (batch_dims, batch_dims)]
+    return lax.dot_general(lhs, rhs, dimension_numbers)
+  else:
+    # move contracting dimensions to the end. lax.dot_general only allows one
+    # contracting dimension, so if there's more than one we collapse them.
+    if ncont > 1:
+      lhs_cdims = tuple(range(lhs.ndim - ncont, lhs.ndim))
+      lhs = moveaxis(lhs, lhs_cont, lhs_cdims).reshape(lhs.shape[:-ncont] + (-1,))
+
+      rhs_cdims = tuple(range(rhs.ndim - ncont, rhs.ndim))
+      rhs = moveaxis(rhs, rhs_cont, rhs_cdims).reshape(rhs.shape[:-ncont] + (-1,))
+    else:
+      lhs = moveaxis(lhs, lhs_cont[0], -1)
+      rhs = moveaxis(rhs, rhs_cont[0], -1)
+
+    # lax.dot_general only allows zero or one tensor product dims per operand,
+    # so if there's more than one we collapse them.
+    result_shape = lhs.shape[:nbatch] + lhs.shape[nbatch:-1] + rhs.shape[nbatch:-1]
+
+    if lhs_ntensor > 1:
+      lhs = lhs.reshape(lhs.shape[:nbatch] + (-1,) + lhs.shape[-1:])
+
+    if rhs_ntensor > 1:
+      rhs = rhs.reshape(rhs.shape[:nbatch] + (-1,) + rhs.shape[-1:])
+
+    lhs_cont, rhs_cont = [lhs.ndim - 1], [rhs.ndim - 1]
+    dimension_numbers = [(lhs_cont, rhs_cont), (batch_dims, batch_dims)]
+    result = lax.dot_general(lhs, rhs, dimension_numbers)
+    return lax.reshape(result, result_shape)
+
+
 ### Misc
 
 


Test patch:
diff --git a/tests/lax_numpy_einsum_test.py b/tests/lax_numpy_einsum_test.py
new file mode 100644
--- /dev/null
+++ b/tests/lax_numpy_einsum_test.py
@@ -0,0 +1,166 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+from absl.testing import absltest
+
+import jax.numpy as np
+import jax.test_util as jtu
+
+from jax.config import config
+config.parse_flags_with_absl()
+
+
+def rng():
+  return onp.random.RandomState(0)
+
+def check(s, *ops):
+  a = onp.einsum(s, *ops)
+  b = np.einsum(s, *ops)
+  assert onp.allclose(a, b, atol=1e-4, rtol=1e-4)
+
+
+class EinsumTest(jtu.JaxTestCase):
+
+  def test_three_operands_1(self):
+    x = rng().randn(3)
+    y = rng().randn(4)
+    z = rng().randn(5)
+    s = 'i,j,k->ijk'
+    check(s, x, y, z)
+
+  def test_three_operands_2(self):
+    x = rng().randn(3)
+    y = rng().randn(4)
+    z = rng().randn(5)
+    s = 'i,j,k->ijk'
+    check(s, x, y, z)
+
+  def test_two_operands_1(self):
+    x = rng().randn(3, 4)
+    y = rng().randn(4)
+    s = 'ij,j->i'
+    check(s, x, y)
+
+  def test_two_operands_2(self):
+    x = rng().randn(3, 4, 5)
+    y = rng().randn(4)
+    s = 'ijk,j->i'
+    check(s, x, y)
+
+  def test_two_operands_3(self):
+    x = rng().randn(3, 4, 3)
+    y = rng().randn(3)
+    s = 'iji,i->j'
+    check(s, x, y)
+
+  def test_two_operands_4(self):
+    x = rng().randn(3, 4)
+    y = rng().randn(3, 4)
+    s = 'ij,ij->'
+    check(s, x, y)
+
+  def test_two_operands_5(self):
+    x = rng().randn(10, 2, 3)
+    y = rng().randn(3, 4)
+    s = 'nij,jk->nik'
+    check(s, x, y)
+
+  def test_one_operand_1(self):
+    x = rng().randn(3, 4, 5)
+    s = 'ijk->j'
+    check(s, x)
+
+  def test_one_operand_2(self):
+    x = rng().randn(3, 4, 5)
+    s = 'ijk->kij'
+    check(s, x)
+
+  def test_one_operand_3(self):
+    x = rng().randn(3, 4, 5)
+    s = 'ijk->ki'
+    check(s, x)
+
+  def test_one_operand_4(self):
+    x = rng().randn(3, 4, 5)
+    s = 'ijk->ki'
+    check(s, x)
+
+  def test_one_operand_5(self):
+    x = rng().randn(2, 3, 4, 5)
+    s = '...ijk->...ki'
+    check(s, x)
+
+  def test_one_operand_6(self):
+    x = rng().randn(3, 4, 5)
+    s = '...ijk->ki'
+    check(s, x)
+
+  def test_one_operand_7(self):
+    x = rng().randn(3, 3)
+    s = 'ii->'
+    check(s, x)
+
+  def test_one_operand_8(self):
+    x = rng().randn(3, 3)
+    s = 'ij->'
+    check(s, x)
+
+  def test_one_operand_9(self):
+    x = rng().randn(3, 3, 3)
+    s = 'iii->'
+    check(s, x)
+
+  def test_one_operand_10(self):
+    x = rng().randn(3, 3)
+    s = 'ii->i'
+    check(s, x)
+
+  def test_one_operand_11(self):
+    x = rng().randn(3, 3, 4)
+    s = 'iij->i'
+    check(s, x)
+
+  def test_one_operand_12(self):
+    x = rng().randn(3, 3, 3)
+    s = 'iii->i'
+    check(s, x)
+
+  def test_one_operand_13(self):
+    x = rng().randn(3, 3, 5, 4, 4)
+    s = 'iijkk->i'
+    check(s, x)
+
+  def test_one_operand_14(self):
+    x = rng().randn(3, 3, 5, 4, 4)
+    s = 'iijkk->ik'
+    check(s, x)
+
+  def test_one_operand_15(self):
+    x = rng().randn(3, 3, 5, 4, 4)
+    s = 'iijkl->il'
+    check(s, x)
+
+  def test_one_operand_16(self):
+    x = rng().randn(3, 3)
+    s = 'ij->ij'
+    check(s, x)
+
+
+if __name__ == '__main__':
+  absltest.main()