Repository: google/jax

Issue:
np.tensordot crashes when axes=0
```python
import jax.numpy as np
import numpy as onp

result = onp.tensordot(onp.ones((2, 3, 4)), onp.ones((5, 6, 7)), 0) 
print(result.shape) # should print (2, 3, 4, 5, 6, 7)
# This errors out.
result = np.tensordot(np.ones((2, 3, 4)), np.ones((5, 6, 7)), 0) 
```
This is the error message I receive.
```
>>> result = np.tensordot(np.ones((2, 3, 4)), np.ones((5, 6, 7)), 0) 
/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 70, in primitive_computation
    return c.Build()
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lib/xla_bridge.py", line 267, in Build
    *args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jaxlib/xla_client.py", line 640, in Build
    return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Cannot infer shape for dot operation: f32[1,24] <dot> f32[1,210]. Contracting dimension sizes do not match.: 

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/numpy/lax_numpy.py", line 1598, in tensordot
    out_reshape = lax.dot(a_reshape, b_reshape)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 462, in dot
    return dot_p.bind(lhs, rhs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/core.py", line 117, in bind
    return self.impl(*args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 51, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 57, in xla_primitive_callable
    built_c = primitive_computation(prim, *shapes, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 73, in primitive_computation
    prim.abstract_eval(*map(aval_from_xla_shape, shapes), **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 1295, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 1824, in _dot_shape_rule
    require(lhs.shape[1] == rhs.shape[0])
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 1818, in require
    raise TypeError(msg.format(lhs.shape, rhs.shape))
TypeError: Incompatible shapes for dot: got (1, 24) and (1, 210).
```


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
@@ -1585,11 +1585,16 @@ def tensordot(a, b, axes=2):
     raise TypeError(msg.format(ndim(a), ndim(b)))
 
   if type(axes) is int:
-    a, b = _promote_dtypes(a, b)
-    a_reshape = lax.reshape(a, (_prod(a.shape[:-axes]), _prod(a.shape[-axes:])))
-    b_reshape = lax.reshape(b, (_prod(b.shape[:axes]), _prod(b.shape[axes:])))
-    out_reshape = lax.dot(a_reshape, b_reshape)
-    return lax.reshape(out_reshape, a.shape[:-axes] + b.shape[axes:])
+    if axes == 0:
+      a, b = _promote_dtypes(a, b)
+      return lax.mul(lax.reshape(a, shape(a) + (1,) * ndim(b)),
+                     lax.reshape(b, (1,) * ndim(a) + shape(b)))
+    else:
+      a, b = _promote_dtypes(a, b)
+      a_reshape = lax.reshape(a, (_prod(a.shape[:-axes]), _prod(a.shape[-axes:])))
+      b_reshape = lax.reshape(b, (_prod(b.shape[:axes]), _prod(b.shape[axes:])))
+      out_reshape = lax.dot(a_reshape, b_reshape)
+      return lax.reshape(out_reshape, a.shape[:-axes] + b.shape[axes:])
   elif type(axes) in (list, tuple) and len(axes) == 2:
     ax1, ax2 = axes
     if type(ax1) == type(ax2) == int:


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -539,6 +539,7 @@ def testMatmul(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, rng):
        "axes": axes, "rng": rng}
       for rng in [jtu.rand_default()]
       for lhs_shape, rhs_shape, axes in [
+          [(2, 3, 4), (5, 6, 7), 0],  # from issue #740
           [(2, 3, 4), (3, 4, 5, 6), 2],
           [(2, 3, 4), (5, 4, 3, 6), [1, 2]],
           [(2, 3, 4), (5, 4, 3, 6), [[1, 2], [2, 1]]],