# Buggy code
def odeint(ofunc, y0, t, *args, **kwargs):
    """Adaptive stepsize (Dormand-Prince) Runge-Kutta odeint implementation.

    Args:
      ofunc: Function to evaluate `yt = ofunc(y, t, *args)` that
        returns the time derivative of `y`.
      y0: initial value for the state.
      t: Timespan for `ofunc` evaluation like `np.linspace(0., 10., 101)`.
      *args: Additional arguments to `ofunc` beyond y0 and t.
      **kwargs: Two relevant keyword arguments:
        '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:
      Integrated system values at each timepoint.
    """
    rtol = kwargs.get("rtol", 1.4e-8)
    atol = kwargs.get("atol", 1.4e-8)
    mxstep = kwargs.get("mxstep", np.inf)

    @functools.partial(jax.jit, static_argnums=(0,))
    def _fori_body_fun(func, i, val):
        """Internal fori_loop body to interpolate an integral at each timestep."""
        t, cur_y, cur_f, cur_t, dt, last_t, interp_coeff, solution = val
        cur_y, cur_f, cur_t, dt, last_t, interp_coeff, _ = jax.lax.while_loop(
            lambda x: (x[2] < t[i]) & (x[-1] < mxstep),
            functools.partial(_while_body_fun, func),
            (cur_y, cur_f, cur_t, dt, last_t, interp_coeff, 0.0),
        )

        relative_output_time = (t[i] - last_t) / (cur_t - last_t)
        out_x = np.polyval(interp_coeff, relative_output_time)

        return (
            t,
            cur_y,
            cur_f,
            cur_t,
            dt,
            last_t,
            interp_coeff,
            jax.ops.index_update(solution, jax.ops.index[i, :], out_x),
        )

    @functools.partial(jax.jit, static_argnums=(0,))
    def _while_body_fun(func, x):
        """Internal while_loop body to determine interpolation coefficients."""
        cur_y, cur_f, cur_t, dt, last_t, interp_coeff, j = x
        next_t = cur_t + dt
        next_y, next_f, next_y_error, k = runge_kutta_step(
            func, cur_y, cur_f, cur_t, dt
        )
        error_ratios = error_ratio(next_y_error, rtol, atol, cur_y, next_y)
        new_interp_coeff = interp_fit_dopri(cur_y, next_y, k, dt)
        dt = optimal_step_size(dt, error_ratios)

        next_j = j + 1
        new_rav, unravel = ravel_pytree(
            (next_y, next_f, next_t, dt, cur_t, new_interp_coeff, next_j)
        )
        old_rav, _ = ravel_pytree(
            (cur_y, cur_f, cur_t, dt, last_t, interp_coeff, next_j)
        )

        return unravel(np.where(np.all(error_ratios <= 1.0), new_rav, old_rav))

    func = lambda y, t: ofunc(y, t, *args)
    f0 = func(y0, t[0])
    dt = initial_step_size(func, t[0], y0, 4, rtol, atol, f0)
    interp_coeff = np.array([y0] * 5)

    return jax.lax.fori_loop(
        1,
        t.shape[0],
        functools.partial(_fori_body_fun, func),
        (
            t,
            y0,
            f0,
            t[0],
            dt,
            t[0],
            interp_coeff,
            jax.ops.index_update(
                np.zeros((t.shape[0], y0.shape[0])), jax.ops.index[0, :], y0
            ),
        ),
    )[-1]


# 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 odeint(ofunc, y0, t, *args, **kwargs):
    """Adaptive stepsize (Dormand-Prince) Runge-Kutta odeint implementation.

    Args:
      ofunc: Function to evaluate `yt = ofunc(y, t, *args)` that
        returns the time derivative of `y`.
      y0: initial value for the state.
      t: Timespan for `ofunc` evaluation like `np.linspace(0., 10., 101)`.
      *args: Additional arguments to `ofunc` beyond y0 and t.
      **kwargs: Two relevant keyword arguments:
        '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:
      Integrated system values at each timepoint.
    """
    rtol = kwargs.get("rtol", 1.4e-8)
    atol = kwargs.get("atol", 1.4e-8)
    mxstep = kwargs.get("mxstep", np.inf)

    func = lambda y, t: ofunc(y, t, *args)

    def _fori_body_fun(i, val):
        """Internal fori_loop body to interpolate an integral at each timestep."""
        t, cur_y, cur_f, cur_t, dt, last_t, interp_coeff, solution = val
        cur_y, cur_f, cur_t, dt, last_t, interp_coeff, _ = jax.lax.while_loop(
            lambda x: (x[2] < t[i]) & (x[-1] < mxstep),
            _while_body_fun,
            (cur_y, cur_f, cur_t, dt, last_t, interp_coeff, 0.0),
        )

        relative_output_time = (t[i] - last_t) / (cur_t - last_t)
        out_x = np.polyval(interp_coeff, relative_output_time)
        solution = jax.ops.index_update(solution, jax.ops.index[i, :], out_x)
        return (t, cur_y, cur_f, cur_t, dt, last_t, interp_coeff, solution)

    def _while_body_fun(x):
        """Internal while_loop body to determine interpolation coefficients."""
        cur_y, cur_f, cur_t, dt, last_t, interp_coeff, j = x
        next_t = cur_t + dt
        next_y, next_f, next_y_error, k = runge_kutta_step(
            func, cur_y, cur_f, cur_t, dt
        )
        error_ratios = error_ratio(next_y_error, rtol, atol, cur_y, next_y)
        new_interp_coeff = interp_fit_dopri(cur_y, next_y, k, dt)
        dt = optimal_step_size(dt, error_ratios)

        next_j = j + 1
        new = (next_y, next_f, next_t, dt, cur_t, new_interp_coeff, next_j)
        old = (cur_y, cur_f, cur_t, dt, last_t, interp_coeff, next_j)
        return tuple(map(partial(np.where, np.all(error_ratios <= 1.0)), new, old))

    f0 = func(y0, t[0])
    dt = initial_step_size(func, t[0], y0, 4, rtol, atol, f0)
    interp_coeff = np.array([y0] * 5)
    solution = jax.ops.index_update(
        np.zeros((t.shape[0], y0.shape[0])), jax.ops.index[0, :], y0
    )
    init_carry = (t, y0, f0, t[0], dt, t[0], interp_coeff, solution)
    *_, solution = jax.lax.fori_loop(1, t.shape[0], _fori_body_fun, init_carry)
    return solution