Repository: jax-ml/jax
Issue category: Feature Request

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]
``




Functions to edit:
jax/_src/custom_derivatives.py:_flatten_bwd

Patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 2752604444dc..9a0a46e89b4b 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -734,12 +734,17 @@ def _flatten_bwd(in_tree, in_avals, out_trees, *args):
   zero = object()  # non-pytree sentinel to replace Nones in py_cts_in
   dummy = tree_unflatten(in_tree, [object()] * in_tree.num_leaves)
   cts_in_flat = []
-  append = lambda x, d: cts_in_flat.extend([x] * len(tree_flatten(d)[0])) or x
+  def append(x, d):
+    num_leaves = len(tree_flatten(d)[0])
+    if x is None and d is not None:
+      cts_in_flat.extend([zero] * num_leaves)
+    elif x is not None:
+      cts_in_flat.extend([x] * num_leaves)
+    return x
   try:
     if not isinstance(py_cts_in, tuple):
       raise ValueError
-    tree_map(append,
-             tuple(zero if ct is None else ct for ct in py_cts_in), dummy)
+    tree_map(append, py_cts_in, dummy, is_leaf=lambda x: x is None)
   except ValueError:
     _, in_tree2 = tree_flatten(py_cts_in)
     msg = ("Custom VJP rule must produce an output with the same container "

Test patch: diff --git a/tests/api_test.py b/tests/api_test.py
index 81b04d5849a7..9ac72ea3bb2b 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9212,6 +9212,38 @@ def g(x):
 
     g(1.)  # doesn't crash
 
+  def test_nones_representing_zeros_in_subtrees_returned_by_bwd(self):
+    # https://github.com/google/jax/issues/8356
+    @jax.custom_vjp
+    def f(x):
+      return x[0]
+
+    def f_fwd(x):
+      return f(x), None
+
+    def f_bwd(_, z_bar):
+      return (z_bar, (None, None)),
+
+    f.defvjp(f_fwd, f_bwd)
+
+    jax.grad(f)((1.0, (2.0, 3.0)))  # don't crash
+
+  def test_pytree_nones_returned_by_bwd(self):
+    @jax.custom_vjp
+    def f(x):
+      return x[0]
+
+    def f_fwd(x):
+      return f(x), None
+
+    def f_bwd(_, z_bar):
+      return (z_bar, (None, None)),
+
+    f.defvjp(f_fwd, f_bwd)
+
+    jax.grad(f)((1.0, (2.0, None)))  # don't crash
+
+
 
 def transpose_unary(f, x_example):
   def transposed(y):