Repository: google/jax

Issue:
`tree_flatten` gets `None` when used in `defvjp`
Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:


I wanted to define a `custom_vjp` for a function that takes a string to indicate what operations are to be conducted. Since `str` is not a validate type for JAX whereas any objects that can be flatten and unflatten as a `pytree` is, I thought of this hacky way to define my own string object: 

```python
import jax
import jax.numpy as jnp
from functools import partial

OP_TO_IDX, IDX_TO_OP = {"add": 0, "mul": 1}, {0: "add", 1: "mul"}

@jax.tree_util.register_pytree_node_class
class OpStr(object):
    def __init__(self, op):
        assert op in OP_TO_IDX, "can only be one of the ops"
        self.op = op

    def __eq__(self, other):
        return self.op == other

    def __hash__(self):
        return hash(self.op)

    def tree_flatten(self):
        return ([OP_TO_IDX[self.op]]), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        idx = int(children[0])
        return cls(op=IDX_TO_OP[idx])

@partial(jax.custom_vjp, nondiff_argnums=(2,))
def op(x, y, op):
    if op == "add": z = x + y
    if op == "mul": z = x * y
    return z

def op_fwd(x, y, op):
    cache = (x, y, op)
    if op == "add": z = x + y
    if op == "mul": z = x * y
    return z, cache

def op_bwd(cache, dz):
    x, y, op = cache
    if op == "add": dz_dx = dz_dy = 1.0
    if op == "mul": dz_dx, dz_dy = y, x
    return dz_dx, dz_dy

op.defvjp(op_fwd, op_bwd)
grad_op = jax.grad(op)
grad_op(1.0, 1.0, OpStr("add"))
```

- [x] If applicable, include full error messages/tracebacks.

But this gives me 
```
Traceback (most recent call last):
  File "/Users/wangy1/Documents/GitHub/dgl/tests/jax/test_op_str.py", line 47, in <module>
    grad_op(1.0, 1.0, OpStr("add"))
  File "/Users/wangy1/Documents/GitHub/dgl/tests/jax/test_op_str.py", line 24, in tree_unflatten
    idx = int(children[0])
TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'
```

Which is not expected since the flattened tree shouldn't be `None`.



Discussion:
It sounds like you're running into the issues discussed here: https://jax.readthedocs.io/en/latest/pytrees.html#custom-pytrees-and-initialization

Do the recommendations there answer your question?
Indeed, to underscore Jake's point, pytrees are technically required to be able to contain any Python object, just like Python tuples can. That is, pytrees are defined in terms of isomorphisms to tuples. Some bits of JAX internals code rely on that property. In this case, we're tree-mapping an assertion function which returns None (which is how Python represents 'no return value'); there's still a result pytree being built, with Nones substituted in place for what were previously leaves. But since this custom pytree node class doesn't handle containing nones, it breaks.

That said, while this bit of JAX internals code is relying on something that technically is part of the pytree contract, usually it's possible for us to rewrite the code not to demand this property hold true (i.e. not to require a custom pytree type to contain Nones). I think we can do that in this case too...

Resolving patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -22,8 +22,9 @@
 from jax import core
 from jax import linear_util as lu
 from jax.custom_transpose import custom_transpose
-from jax.tree_util import (tree_flatten, tree_unflatten, tree_map, treedef_is_leaf,
-                           treedef_tuple, register_pytree_node_class)
+from jax.tree_util import (tree_flatten, tree_unflatten, tree_map,
+                           treedef_is_leaf, treedef_tuple,
+                           register_pytree_node_class, tree_leaves)
 from jax._src import custom_api_util
 from jax._src import dtypes
 from jax._src.util import cache, safe_zip, safe_map, split_list, Unhashable
@@ -564,21 +565,17 @@ def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:  # pytype: disable
       out_tree = aux if fst else aux[0]
       return tree_unflatten(out_tree, out_flat)
 
-@partial(partial, tree_map)
 def _check_for_tracers(x):
-  if isinstance(x, core.Tracer):
-    msg = ("Found a JAX Tracer object passed as an argument to a custom_vjp "
-           "function in a position indicated by nondiff_argnums as "
-           "non-differentiable. Tracers cannot be passed as non-differentiable "
-           "arguments to custom_vjp functions; instead, nondiff_argnums should "
-           "only be used for arguments that can't be or contain JAX tracers, "
-           "e.g. function-valued arguments. In particular, array-valued "
-           "arguments should typically not be indicated as nondiff_argnums. "
-           "\n\n"
-           "This behavior recently changed in JAX. "
-           "See https://github.com/google/jax/blob/main/docs/custom_vjp_update.md "
-           "for more information.")
-    raise UnexpectedTracerError(msg)
+  for leaf in tree_leaves(x):
+    if isinstance(x, core.Tracer):
+      msg = ("Found a JAX Tracer object passed as an argument to a custom_vjp "
+            "function in a position indicated by nondiff_argnums as "
+            "non-differentiable. Tracers cannot be passed as non-differentiable "
+            "arguments to custom_vjp functions; instead, nondiff_argnums should "
+            "only be used for arguments that can't be or contain JAX tracers, "
+            "e.g. function-valued arguments. In particular, array-valued "
+            "arguments should typically not be indicated as nondiff_argnums.")
+      raise UnexpectedTracerError(msg)
 
 @lu.transformation_with_aux
 def _flatten_fwd(in_tree, *args):
@@ -610,11 +607,11 @@ 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_cts = lambda x, d: cts_in_flat.extend([x] * len(tree_flatten(d)[0]))
+  append = lambda x, d: cts_in_flat.extend([x] * len(tree_flatten(d)[0])) or x
   try:
     if not isinstance(py_cts_in, tuple):
       raise ValueError
-    tree_map(append_cts,
+    tree_map(append,
              tuple(zero if ct is None else ct for ct in py_cts_in), dummy)
   except ValueError:
     _, in_tree2 = tree_flatten(py_cts_in)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -6442,6 +6442,26 @@ def g_rev(res, w_bar):
 
     jtu.check_grads(h, (jnp.float32(3.14),), order=1, modes=['rev'])
 
+  def test_pytrees_not_required_to_contain_nones(self):
+    class A(list):
+      pass
+
+    def unflatten(_, children):
+      assert children[0] is not None
+      return A(children)
+
+    tree_util.register_pytree_node(A, lambda x: (x, None), unflatten)
+
+    @jax.custom_vjp
+    def f(x):
+      return x[0]
+    def f_fwd(x):
+      return x[0], None
+    def f_bwd(_, g):
+      return A([g]),
+    f.defvjp(f_fwd, f_bwd)
+
+    jax.grad(f)(A([1.]))  # doesn't crash
 
 def transpose_unary(f, x_example):
   def transposed(y):