Improve error message in custom_vjp  when returning the wrong grad structure
```
import jax
from jax import numpy as jnp


def splat_core(data, warp, output_dims):
  output = jnp.zeros((output_dims[0], output_dims[1], data.shape[-1]),
                     dtype=data.dtype)
  output = output.at[:data.shape[0], :data.shape[1]].set(data * warp[..., :1] +
                                                         warp[..., 1:2])
  return output


@jax.custom_vjp
def splat_2d(data, warp, output_dims):
  """Functions.

  Args:
    data: An [H, W, C] array.
    warp: An [H, W, 2] array.
    output_dims: A 2-tuple containing the output height and width.

  Returns:
    An array with shape [output_dims[0], output_dims[1], data.shape[2]]
  """
  return splat_core(data, warp, output_dims=output_dims)


def splat_2d_fwd(data, warp, output_dims):
  return splat_2d(data, warp, output_dims), (data, warp)


def splat_2d_grad_impl(res, g):
  return res[0] * 2.0 * jax.numpy.sum(g), res[1] * 2.0 * jax.numpy.sum(g)


def splat_2d_bwd(res, g):
  actual_grad = res[0] * 2.0 * jax.numpy.sum(g), res[1] * 2.0 * jax.numpy.sum(g)
  # Need to return a tuple of None's here since output_dims is intepreted as
  # a pytree.
  return actual_grad + ((None, None),)


splat_2d.defvjp(splat_2d_fwd, splat_2d_bwd)


def test():
  output_dims = (6, 7)
  warp = jnp.ones((5, 5, 2))
  data = jnp.ones((5, 5, 3))
  grad_output = jnp.ones((output_dims[0], output_dims[1], 3))
  # Is passing the output_dims here right?
  actual_grad = jax.vjp(splat_2d, data, warp, output_dims)[1](grad_output)

test()
```
Note that `splat_2d_bwd` should return actual_grad + (None,)
Gives this confusing error message:
```
The above exception was the direct cause of the following exception:
 
AssertionError                            Traceback (most recent call last)
<ipython-input-5-e07748477af3> in <module>()
     50 
     51 
---> 52 test()
 
<ipython-input-5-e07748477af3> in test()
     47   data = jnp.ones((5, 5, 3))
     48   grad_output = jnp.ones((output_dims[0], output_dims[1], 3))
---> 49   actual_grad = jax.vjp(splat_2d, data, warp, output_dims)[1](grad_output)
     50 
     51 
 
google3/third_party/py/jax/_src/tree_util.py in <lambda>(*args, **kw)
    324     if isinstance(func, functools.partial):
    325       original_func = func
--> 326       func = lambda *args, **kw: original_func(*args, **kw)
    327       func.func = original_func.func
    328       func.args = original_func.args
 
google3/third_party/py/jax/_src/api.py in _vjp_pullback_wrapper(cotangent_dtypes, cotangent_shapes, io_tree, fun, py_args)
   2217           "must be the same as the shape of corresponding primal input "
   2218           f"{ct_shape}.")
-> 2219   ans = fun(*args)
   2220   return tree_unflatten(out_tree, ans)
   2221 
 
google3/third_party/py/jax/_src/tree_util.py in <lambda>(*args, **kw)
    324     if isinstance(func, functools.partial):
    325       original_func = func
--> 326       func = lambda *args, **kw: original_func(*args, **kw)
    327       func.func = original_func.func
    328       func.args = original_func.args
 
google3/third_party/py/jax/interpreters/ad.py in unbound_vjp(pvals, jaxpr, consts, *cts)
    121     cts = tuple(map(ignore_consts, cts, pvals))
    122     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
--> 123     arg_cts = backward_pass(jaxpr, reduce_axes, consts, dummy_args, cts)
    124     return map(instantiate_zeros, arg_cts)
    125 
 
google3/third_party/py/jax/interpreters/ad.py in backward_pass(jaxpr, reduce_axes, consts, primals_in, cotangents_in)
    227       else:
    228         cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
--> 229                                                          **eqn.params)
    230     cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out
    231     # FIXME: Some invars correspond to primals!
 
google3/third_party/py/jax/interpreters/ad.py in _custom_lin_transpose(cts_out, num_res, bwd, out_avals, *invals)
    688   res, _ = split_list(invals, [num_res])
    689   cts_out = map(instantiate_zeros_aval, out_avals, cts_out)
--> 690   cts_in = bwd.call_wrapped(*res, *cts_out)
    691   return [None] * num_res + list(cts_in)
    692 primitive_transposes[custom_lin_p] = _custom_lin_transpose
 
google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    177     while stack:
    178       gen, out_store = stack.pop()
--> 179       ans = gen.send(ans)
    180       if out_store is not None:
    181         ans, side = ans
 
google3/third_party/py/jax/_src/custom_derivatives.py in _flatten_bwd(in_tree, in_avals, out_trees, *args)
    592   # TODO(mattjj): change this to check if tangent type represents 0dim vspace
    593   yield [Zero(a.at_least_vspace()) if ct is zero or a != a.at_least_vspace()
--> 594          else ct for a, ct in zip(in_avals, cts_in_flat)]
    595 
    596 
 
google3/third_party/py/jax/_src/util.py in safe_zip(*args)
     33   n = len(args[0])
     34   for arg in args[1:]:
---> 35     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
     36   return list(zip(*args))
     37 
 
AssertionError: length mismatch: [4, 2]
``