improve error message with when custom_vjp bwd rule produces wrong shape/dtype

[mattjj]

Even the best of us can miss important error information, like noticing when shapes disagree. So let's make things better.

import jax
import jax.numpy as jnp

@jax.custom_vjp
def f(x):
  return x

def f_fwd(x):
  return x, ()

def f_bwd(_, g):
  return jnp.ones((10,), g.dtype),

f.defvjp(f_fwd, f_bwd)

jax.grad(f)(3.14)

Before:
Traceback (most recent call last):
File "/usr/local/google/home/mattjj/packages/jax/rahul11.py", line 16, in
jax.grad(f)(3.14)
ValueError: Custom VJP bwd rule must produce an output with the same shape/dtypes as the args tuple of the primal function, but at output[0] the bwd rule produced an output of shape/dtype float32[10] corresponding to an input of shape/dtype float32[].

After (emphasis added):
Traceback (most recent call last):
File "/usr/local/google/home/mattjj/packages/jax/rahul11.py", line 16, in
jax.grad(f)(3.14)
ValueError: Custom VJP bwd rule must produce an output with the same shape/dtypes as the args tuple of the primal function, but at output[0] the bwd rule produced an output of shape/dtype float32[10] corresponding to an input of shape/dtype float32[], so the shapes do not match

[froystig]

Looks good (hindsight review).