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.
Please reopen if you notice any missing cases!
thanks @mattjj 

I tested with some existing code for which I'd like to use jax

```
import jax.numpy as np
import numpy as onp

second = onp.random.uniform(size = (2,3,4))
first = onp.random.uniform(size = (2,1))
result = np.einsum('sa,shb->shab',first,second)
result.shape #(expected shape: (2, 3, 1, 4))
```

which throws a `StopIteration` error, while changing `import jax.numpy as np` to `import numpy as np
` gives the expected shape
Sounds like a bug! 

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
@@ -1155,14 +1155,16 @@ def sum_repeats(operand, names, counts, keep_names):
         if lhs_batch != rhs_batch:
           rhs = moveaxis(rhs, rhs_batch, lhs_batch)
         batch_names = ''.join(lhs_names[i] for i in lhs_batch)
+        nbatch = len(batch_names)
 
-        names = batch_names + lhs_names + rhs_names
+        assert len(lhs_names) == lhs.ndim and len(rhs_names) == rhs.ndim
+        assert lhs_names.startswith(batch_names) and rhs_names.startswith(batch_names)
+
+        names = batch_names + lhs_names[nbatch:] + rhs_names[nbatch:]
         lhs_shape = iter(lhs.shape)
-        lhs_shape = [next(lhs_shape) if n in batch_names + lhs_names else 1
-                     for n in names]
+        lhs_shape = [next(lhs_shape) if n in 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]
+        rhs_shape = [next(rhs_shape) if n in rhs_names else 1 for n in names]
         operand = lax.reshape(lhs, lhs_shape) * lax.reshape(rhs, rhs_shape)
 
     else:


Test patch:
diff --git a/tests/lax_numpy_einsum_test.py b/tests/lax_numpy_einsum_test.py
--- a/tests/lax_numpy_einsum_test.py
+++ b/tests/lax_numpy_einsum_test.py
@@ -81,6 +81,13 @@ def test_two_operands_5(self):
     s = 'nij,jk->nik'
     check(s, x, y)
 
+  def test_two_operands_6(self):
+    # based on https://github.com/google/jax/issues/37#issuecomment-448572187
+    x = rng().randn(2, 1)
+    y = rng().randn(2, 3, 4)
+    s = 'sa,shb->shab'
+    check(s, x, y)
+
   def test_one_operand_1(self):
     x = rng().randn(3, 4, 5)
     s = 'ijk->j'