Repository: google/jax

Issue:
einsum broadcasting failure
```python
import jax.numpy as jnp
import numpy as np
x = np.zeros((2, 1, 3, 3))
y = np.zeros((2, 4, 3))
assert np.einsum('...ij,...j', x, y).shape == (2, 4, 3)  # passes
assert jnp.einsum('...ij,...j', x, y).shape == (2, 4, 3)  # raises an error
```
Errors with:
`TypeError: dot_general requires lhs batch dimensions and rhs batch dimensions to have the same shape, got [2 4] and [2 1].`

Full traceback:
<details>

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-ba41a135a1c7> in <module>()
      5 y = np.zeros((2, 4, 3))
      6 assert np.einsum('...ij,...j', x, y).shape == (2, 4, 3)
----> 7 assert jnp.einsum('...ij,...j', x, y).shape == (2, 4, 3)
      8 # jnp.einsum('...ij,...j', x, y)
/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in einsum(*operands, **kwargs)
   2342       *operands, einsum_call=True, use_blas=True, optimize=optimize)
   2343   contractions = tuple(data[:3] for data in contractions)
-> 2344   return _einsum(operands, contractions, precision)
   2345 
   2346 @_wraps(onp.einsum_path)

/usr/local/lib/python3.6/dist-packages/jax/api.py in f_jitted(*args, **kwargs)
    148     flat_fun, out_tree = flatten_fun(f, in_tree)
    149     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend,
--> 150                        name=flat_fun.__name__)
    151     return tree_unflatten(out_tree(), out)
    152 

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, f, *args, **params)
    603   if top_trace is None:
    604     with new_sublevel():
--> 605       outs = primitive.impl(f, *args, **params)
    606   else:
    607     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, *args)
    447 
    448 def _xla_call_impl(fun, *args, device, backend, name):
--> 449   compiled_fun = _xla_callable(fun, device, backend, name, *map(arg_spec, args))
    450   try:
    451     return compiled_fun(*args)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
    221       fun.populate_stores(stores)
    222     else:
--> 223       ans = call(fun, *args)
    224       cache[key] = (ans, fun.stores)
    225     return ans

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, *arg_specs)
    464   pvals = [pe.PartialVal((aval, core.unit)) for aval in abstract_args]
    465   with core.new_master(pe.StagingJaxprTrace, True) as master:
--> 466     jaxpr, (pvals, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
    467     assert not env  # no subtraces here
    468     del master, env

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    150     gen = None
    151 
--> 152     ans = self.f(*args, **dict(self.params, **kwargs))
    153     del args
    154     while stack:

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in _einsum(operands, contractions, precision)
   2444       bdims = tuple(range(len(batch_dims)))
   2445       dimension_numbers = [(lhs_cont, rhs_cont), (bdims, bdims)]
-> 2446       operand = lax.dot_general(lhs, rhs, dimension_numbers, precision)
   2447       deleted_names = batch_names + ''.join(contracted_names)
   2448       names = (batch_names + removechars(lhs_names, deleted_names)

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in dot_general(lhs, rhs, dimension_numbers, precision)
    615   return dot_general_p.bind(lhs, rhs,
    616                             dimension_numbers=(contract_dims, batch_dims),
--> 617                             precision=_canonicalize_precision(precision))
    618 
    619 def broadcast(operand, sizes):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    160 
    161     tracers = map(top_trace.full_raise, args)
--> 162     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    163     if self.multiple_results:
    164       return map(full_lower, out_tracer)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
    104       tracers = map(self.instantiate_const, tracers)
    105       avals = [t.aval for t in tracers]
--> 106       out_aval = primitive.abstract_eval(*avals, **params)
    107       if primitive.multiple_results:
    108         out_tracers = [JaxprTracer(self, PartialVal((aval, unit)), None)

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in standard_abstract_eval(prim, shape_rule, dtype_rule, *args, **kwargs)
   1540     return ConcreteArray(prim.impl(*[x.val for x in args], **kwargs))
   1541   elif least_specialized is ShapedArray:
-> 1542     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1543   elif least_specialized is UnshapedArray:
   1544     return UnshapedArray(dtype_rule(*args, **kwargs))

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _dot_general_shape_rule(lhs, rhs, dimension_numbers, precision)
   2186     msg = ("dot_general requires lhs batch dimensions and rhs batch dimensions "
   2187            "to have the same shape, got {} and {}.")
-> 2188     raise TypeError(msg.format(lhs_batch_shape, rhs_batch_shape))
   2189   if tuple(sorted(lhs_batch)) != tuple(range(len(lhs_batch))):
   2190     msg = ("dot_general requires lhs batch dimensions to precede contracting "

TypeError: dot_general requires lhs batch dimensions and rhs batch dimensions to have the same shape, got [2 4] and [2 1].
```
</details>


Discussion:

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
@@ -2317,25 +2317,20 @@ def dot(a, b, precision=None):  # pylint: disable=missing-docstring
 def matmul(a, b, precision=None):  # pylint: disable=missing-docstring
   _check_arraylike("matmul", a, b)
   a_is_vec, b_is_vec = (ndim(a) == 1), (ndim(b) == 1)
-  a = lax.reshape(a, (1,) + shape(a)) if a_is_vec else a
-  b = lax.reshape(b, shape(b) + (1,)) if b_is_vec else b
-
-  a, b = _promote_dtypes(a, b)
-  batch_shape = lax.broadcast_shapes(shape(a)[:-2], shape(b)[:-2])
-  a = broadcast_to(a, batch_shape + shape(a)[-2:])
-  b = broadcast_to(b, batch_shape + shape(b)[-2:])
-  batch_dims = tuple(range(len(batch_shape)))
-  dim_numbers = (((ndim(a) - 1,), (ndim(b) - 2,)), (batch_dims, batch_dims))
-  result = lax.dot_general(a, b, dim_numbers,  precision)
-
-  if a_is_vec or b_is_vec:
-    m, n = shape(result)[-2:]
-    new_m = () if a_is_vec else (m,)
-    new_n = () if b_is_vec else (n,)
-    return lax.reshape(result, batch_shape + new_m + new_n)
+  # We lower to einsum here because it handles batch dimensions for us.
+  # np.matmul is stricter than np.einsum with respect to size 1 contracting
+  # dimensions, so we need an additional check.
+  if shape(a)[0 if a_is_vec else -1] != shape(b)[0 if b_is_vec else -2]:
+    msg = "matmul requires contracting dimension to match, got {} and {}"
+    raise ValueError(msg.format(shape(a), shape(b)))
+  if a_is_vec and b_is_vec:
+    return lax.dot(a, b, precision=precision)
+  elif a_is_vec:
+    return einsum('i,...ij->...j', a, b, precision=precision)
+  elif b_is_vec:
+    return einsum('...ij,j->...i', a, b, precision=precision)
   else:
-    return result
-
+    return einsum('...ij,...jk->...ik', a, b, precision=precision)
 
 @_wraps(onp.vdot, lax_description=_PRECISION_DOC)
 def vdot(a, b, precision=None):
@@ -2425,6 +2420,17 @@ def sum_repeats(operand, names, counts, keep_names):
           names = names.replace(name, '', count - 1)
     return operand, names
 
+  def filter_singleton_dims(operand, names, other_shape, other_names):
+    s = shape(operand)
+    new_shape = []
+    new_names = []
+    for i, d in enumerate(names):
+      other_i = other_names.find(d)
+      if s[i] != 1 or other_i == -1 or other_shape[other_i] == 1:
+        new_shape.append(s[i])
+        new_names.append(d)
+    return reshape(operand, tuple(new_shape)), "".join(new_names)
+
   for operand_indices, contracted_names, einstr in contractions:
     input_str, result_names = einstr.split('->')
     input_names = input_str.split(',')
@@ -2445,9 +2451,18 @@ def sum_repeats(operand, names, counts, keep_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
 
+      # handle cases where one side of a contracting or batch dimension is 1
+      # but its counterpart is not.
+      lhs, lhs_names = filter_singleton_dims(lhs, lhs_names, shape(rhs),
+                                             rhs_names)
+      rhs, rhs_names = filter_singleton_dims(rhs, rhs_names, shape(lhs),
+                                             lhs_names)
+
+      lhs_counts = collections.Counter(lhs_names)
+      rhs_counts = collections.Counter(rhs_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]
@@ -2465,14 +2480,16 @@ def sum_repeats(operand, names, counts, keep_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)
 
       # NOTE(mattjj): this can fail non-deterministically in python3, maybe
       # due to opt_einsum
-      assert _all(name in lhs_names and name in rhs_names and
-                  lhs.shape[lhs_names.index(name)] == rhs.shape[rhs_names.index(name)]
-                  for name in contracted_names)
+      assert _all(
+        name in lhs_names and name in rhs_names and
+        lhs.shape[lhs_names.index(name)] == rhs.shape[rhs_names.index(name)]
+        for name in contracted_names)
 
       # move batch dims to the front (required by lax.dot_general, and easier)
       batch_dims = tuple(range(len(batch_names)))


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
@@ -315,6 +315,34 @@ def test_einsum_kpmurphy_example(self):
     self.assertAllClose(L, np.einsum('ntk,kd,dc->nc', S, W, V, optimize=path),
                         check_dtypes=False, rtol=rtol)
 
+  def test_contraction_broadcasting(self):
+    r = rng()
+    x = r.randn(3, 4, 5)
+    y = r.randn(3, 1, 6)
+    s = 'cij,cjk->cik'
+    self._check(s, x, y)
+
+  def test_batch_broadcasting(self):
+    r = rng()
+    x = r.randn(1, 4, 5)
+    y = r.randn(3, 5, 6)
+    s = 'cij,cjk->cik'
+    self._check(s, x, y)
+
+  def test_batch_and_contraction_broadcasting(self):
+    r = rng()
+    x = r.randn(1, 4, 5)
+    y = r.randn(3, 1, 6)
+    s = 'cij,cjk->cik'
+    self._check(s, x, y)
+
+  def test_broadcasting_issue_2189(self):
+    r = rng()
+    x = r.randn(2, 1, 3, 3)
+    y = r.randn(2, 4, 3)
+    s = '...ij,...j'
+    self._check(s, x, y)
+
 
 if __name__ == '__main__':
   absltest.main()