# Buggy code
def tensordot(a, b, axes=2):
    _check_arraylike("tensordot", a, b)
    if not (ndim(a) >= 1 and ndim(b) >= 1):
        msg = "tensordot requires a.ndim and b.dim to be at least 1, got {} and {}."
        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:])
    elif type(axes) in (list, tuple) and len(axes) == 2:
        ax1, ax2 = axes
        if type(ax1) == type(ax2) == int:
            a_transposed = moveaxis(a, ax1, -1) if ax1 != a.ndim - 1 else a
            b_transposed = moveaxis(b, ax2, 0) if ax2 != 0 else b
            return tensordot(a_transposed, b_transposed, 1)
        elif type(ax1) in (list, tuple) and type(ax2) in (list, tuple):
            if len(ax1) != len(ax2):
                msg = (
                    "tensordot requires axes lists to have equal length, got {} and {}."
                )
                raise TypeError(msg.format(ax1, ax2))
            num_axes = len(ax1)
            a_transposed = moveaxis(a, ax1, tuple(range(a.ndim - num_axes, a.ndim)))
            b_transposed = moveaxis(b, ax2, tuple(range(num_axes)))
            return tensordot(a_transposed, b_transposed, num_axes)
    msg = (
        "tensordot axes argument must be an int, a pair of ints, or a pair of "
        "lists/tuples of ints."
    )
    raise TypeError(msg)


# Traceback (RuntimeError)
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).

# Fixed code
def tensordot(a, b, axes=2):
    _check_arraylike("tensordot", a, b)
    if not (ndim(a) >= 1 and ndim(b) >= 1):
        msg = "tensordot requires a.ndim and b.dim to be at least 1, got {} and {}."
        raise TypeError(msg.format(ndim(a), ndim(b)))

    if type(axes) is int:
        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:
            a_transposed = moveaxis(a, ax1, -1) if ax1 != a.ndim - 1 else a
            b_transposed = moveaxis(b, ax2, 0) if ax2 != 0 else b
            return tensordot(a_transposed, b_transposed, 1)
        elif type(ax1) in (list, tuple) and type(ax2) in (list, tuple):
            if len(ax1) != len(ax2):
                msg = (
                    "tensordot requires axes lists to have equal length, got {} and {}."
                )
                raise TypeError(msg.format(ax1, ax2))
            num_axes = len(ax1)
            a_transposed = moveaxis(a, ax1, tuple(range(a.ndim - num_axes, a.ndim)))
            b_transposed = moveaxis(b, ax2, tuple(range(num_axes)))
            return tensordot(a_transposed, b_transposed, num_axes)
    msg = (
        "tensordot axes argument must be an int, a pair of ints, or a pair of "
        "lists/tuples of ints."
    )
    raise TypeError(msg)