Repository: google/jax

Issue:
Shape-checking bug
Hi there,

I've encountered a bug with Jax's shape-checking rules that directed me to report to this forum. I've worked down to roughly the minimum self-contained bug-inducing example:

```python
import jax.numpy as np
from jax import grad, vmap

a = np.arange(4) # example vector
A = np.arange(16).reshape((4, 4)) # example matrix

def fun(x):
    x @ A # this line shouldn't do anything, but commenting it out makes this code run without errors!
    return np.sum(a * x)

jac_diag = vmap(grad(fun)) # element-wise gradient

input = np.arange(4).astype(float) # example input

print(fun(input)) # This runs fine
print(jac_diag(input)) # This throws a "shape-checking" error, see trace below
```
When I run this, I get an error with a long stack trace that I've copied to [Pastebin](https://pastebin.com/5ctS14SU). I'm running on Windows 10 with Windows Subsystem for Linux running Ubuntu 18.04, set up as a remote interpreter through PyCharm.

I suspect it might have something to do with the `@` operator - any ideas here? ~Rewriting that as a series of `np.dot`s might fix the issue~, but I'm curious about what's causing this.


Discussion:
Update: rewriting with `np.dot`s did not resolve this problem.
Thanks for the report! Looks like a bug in `jax.numpy.matmul`, as well as a bug in our shape checking rules (the latter because XLA caught this rather than us catching it earlier in JAX).

However, replacing `x @ A` with `np.dot(x, A)` made the issue go away for me. If it's not working for you, could you share the code that repros it?

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2657,6 +2657,20 @@ def _precision_config(precision):
 
 def _dot_general_shape_rule(lhs, rhs, *, dimension_numbers, precision):
   (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
+  if not all(onp.all(onp.greater_equal(d, 0)) and onp.all(onp.less(d, lhs.ndim))
+             for d in (lhs_contracting, lhs_batch)):
+    msg = ("dot_general requires lhs dimension numbers to be nonnegative and "
+           "less than the number of axes of the lhs value, got "
+           f"lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} "
+           f"for lhs of rank {lhs.ndim}")
+    raise TypeError(msg)
+  if not all(onp.all(onp.greater_equal(d, 0)) and onp.all(onp.less(d, rhs.ndim))
+             for d in (rhs_contracting, rhs_batch)):
+    msg = ("dot_general requires rhs dimension numbers to be nonnegative and "
+           "less than the number of axes of the rhs value, got "
+           f"rhs_batch of {rhs_batch} and rhs_contracting of {rhs_contracting} "
+           f"for rhs of rank {rhs.ndim}")
+    raise TypeError(msg)
   if len(lhs_batch) != len(rhs_batch):
     msg = ("dot_general requires equal numbers of lhs_batch and rhs_batch "
            "dimensions, got lhs_batch {} and rhs_batch {}.")
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
@@ -2783,6 +2783,12 @@ def dot(a, b, *, precision=None):  # pylint: disable=missing-docstring
 @_wraps(np.matmul, lax_description=_PRECISION_DOC)
 def matmul(a, b, *, precision=None):  # pylint: disable=missing-docstring
   _check_arraylike("matmul", a, b)
+  for i, x in enumerate((a, b)):
+    if ndim(x) < 1:
+      msg = (f"matmul input operand {i} must have ndim at least 1, "
+             f"but it has ndim {ndim(x)}")
+      raise ValueError(msg)
+
   a_is_vec, b_is_vec = (ndim(a) == 1), (ndim(b) == 1)
   a = expand_dims(a, axis=0) if a_is_vec else a
   b = expand_dims(b, axis=-1) if b_is_vec else b


Test patch: