# Buggy code
def build_odeint(ofunc, rtol=1.4e-8, atol=1.4e-8, mxstep=onp.inf):
    """Return `f(y0, t, args) = odeint(ofunc(y, t, *args), y0, t, args)`.

    Given the function ofunc(y, t, *args), return the jitted function
    `f(y0, t, args) = odeint(ofunc(y, t, *args), y0, t, args)` with
    the VJP of `f` defined using `vjp_odeint`, where:

      `y0` is the initial condition of the ODE integration,
      `t` is the time course of the integration, and
      `*args` are all other arguments to `ofunc`.

    Args:
      ofunc: The function to be wrapped into an ODE integration.
      rtol: relative local error tolerance for solver.
      atol: absolute local error tolerance for solver.
      mxstep: Maximum number of steps to take for each timepoint.

    Returns:
      `f(y0, t, args) = odeint(ofunc(y, t, *args), y0, t, args)`
    """
    ct_odeint = jax.custom_transforms(
        lambda y0, t, *args: odeint(
            ofunc, y0, t, *args, rtol=rtol, atol=atol, mxstep=mxstep
        )
    )

    v = lambda y0, t, *args: vjp_odeint(
        ofunc, y0, t, *args, rtol=rtol, atol=atol, mxstep=mxstep
    )
    jax.defvjp_all(ct_odeint, v)

    return jax.jit(ct_odeint)


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-91-eb0f4b094a7d> in <module>()
12 args = ([1.0, {}],)
13 print(identity(args))
---> 14 jax.jvp(identity, args, args)

/usr/local/lib/python3.6/dist-packages/jax/api.py in jvp(fun, primals, tangents)
866   ps_flat, ts_flat, in_trees = unzip3(map(trim_arg, primals, tangents))
867   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
--> 868   out_primal, out_tangent = ad.jvp(jaxtree_fun).call_wrapped(ps_flat, ts_flat)
869   return (build_tree(out_tree(), out_primal), build_tree(out_tree(), out_tangent))
870

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
147
148     del gen
--> 149     ans = self.f(*args, **dict(self.params, **kwargs))
150     del args
151     while stack:

/usr/local/lib/python3.6/dist-packages/jax/api.py in __call__(self, *args, **kwargs)
1177     jaxpr, _, consts = pe.trace_to_jaxpr(jaxtree_fun, pvals_in, instantiate=True)
1178     ans = self.prim.bind(core.pack(consts), jax_kwargs, *jax_args,
-> 1179                          in_trees=in_trees, jaxpr=jaxpr)
1180     return build_tree(out_tree(), ans)
1181

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
145
146     tracers = map(top_trace.full_raise, args)
--> 147     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
148     return full_lower(out_tracer)
149

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
250           "Forward-mode differentiation rule for '{}' not implemented"
251           .format(primitive))
--> 252     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
253     return JVPTracer(self, primal_out, tangent_out)
254

/usr/local/lib/python3.6/dist-packages/jax/api.py in custom_transforms_jvp(primals, tangents, **params)
1321     in_trees = params['in_trees']
1322     args = tuple(map(build_tree, in_trees, jax_args))
-> 1323     args_dot = tuple(map(build_tree, in_trees, jax_args_dot))
1324     pytree_out, pytree_out_dot = custom_jvp(args, args_dot)
1325     out, out_tree = pytree_to_jaxtupletree(pytree_out)

/usr/local/lib/python3.6/dist-packages/jax/util.py in safe_map(f, *args)
39   for arg in args[1:]:
40     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
---> 41   return list(map(f, *args))
42
43

/usr/local/lib/python3.6/dist-packages/jax/tree_util.py in build_tree(treedef, xs)
202     else:
203       # We use 'iter' for clearer error messages
--> 204       children = safe_map(build_tree, iter(treedef.children), iter(xs))
205       return treedef.node_type.from_iterable(treedef.node_data, children)
206

/usr/local/lib/python3.6/dist-packages/jax/util.py in safe_map(f, *args)
39   for arg in args[1:]:
40     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
---> 41   return list(map(f, *args))
42
43

/usr/local/lib/python3.6/dist-packages/jax/tree_util.py in build_tree(treedef, xs)
202     else:
203       # We use 'iter' for clearer error messages
--> 204       children = safe_map(build_tree, iter(treedef.children), iter(xs))
205       return treedef.node_type.from_iterable(treedef.node_data, children)
206

TypeError: 'Zero' object is not iterable

# Fixed code
def build_odeint(ofunc, rtol=1.4e-8, atol=1.4e-8, mxstep=np.inf):
    """Return `f(y0, t, args) = odeint(ofunc(y, t, *args), y0, t, args)`.

    Given the function ofunc(y, t, *args), return the jitted function
    `f(y0, t, args) = odeint(ofunc(y, t, *args), y0, t, args)` with
    the VJP of `f` defined using `vjp_odeint`, where:

      `y0` is the initial condition of the ODE integration,
      `t` is the time course of the integration, and
      `*args` are all other arguments to `ofunc`.

    Args:
      ofunc: The function to be wrapped into an ODE integration.
      rtol: relative local error tolerance for solver.
      atol: absolute local error tolerance for solver.
      mxstep: Maximum number of steps to take for each timepoint.

    Returns:
      `f(y0, t, args) = odeint(ofunc(y, t, *args), y0, t, args)`
    """
    fwd = partial(odeint, ofunc, rtol=rtol, atol=atol, mxstep=mxstep)
    bwd = partial(vjp_odeint, ofunc, rtol=rtol, atol=atol, mxstep=mxstep)
    return custom_gradient(fwd, bwd)