Merge pull request #397 from google/fix-dot-batch-rule

fix broken dot batch rule case

jax/lax.py

@@ -1326,11 +1326,9 @@ def _dot_batch_rule(batched_args, batch_dims):
 
     assert lbd is not None and rbd is not None
     assert lhs.ndim == rhs.ndim == 2  # dot only supports rank 1 and above
-    if lbd != 0:
-      batching.move_dim_to_front(lhs, lbd)
-    if rbd != 0:
-      batching.move_dim_to_front(rhs, rbd)
-    return dot_general(lhs, rhs, [((1,), (1,)), ((0,), (0,))])
+    lhs = batching.move_dim_to_front(lhs, lbd)
+    rhs = batching.move_dim_to_front(rhs, rbd)
+    return dot_general(lhs, rhs, [((1,), (1,)), ((0,), (0,))]), 0
 
   if lbd is None:
     assert rbd is not None

tests/batching_test.py

@@ -271,8 +271,15 @@ def vecvec(a, b):
 
     assert vecvec(np.zeros((3,)), np.zeros((3,))).shape == ()
     assert vecvec(np.zeros((2, 3)), np.zeros((3,))).shape == (2,)
-    # TODO(mattjj): this fails due to an xla error in dot_general
-    # assert vecvec(np.zeros((4, 2, 3)), np.zeros((3,))).shape == (4, 2)
+    assert vecvec(np.zeros((4, 2, 3)), np.zeros((3,))).shape == (4, 2)
+
+  def testDot2(self):
+    R = onp.random.RandomState(0).randn
+    xs = R(10, 3)
+    ys = R(10, 3)
+    ans = vmap(np.dot)(xs, ys)
+    expected = onp.einsum('ni,ni->n', xs, ys)
+    self.assertAllClose(ans, expected, check_dtypes=False)
 
   def testPad(self):
     R = onp.random.RandomState(0).randn