human
Complete this python program:
def checkpoint(fun):
assistant
#Here is the Python code you requested.

    """Returns a checkpointed version of `fun`, where intermediate values
    computed during the forward pass of `fun` are discarded and then recomputed
    for the backward pass. Useful to save memory, effectively trading off time
    and memory. See e.g. arxiv.org/abs/1604.06174.
    """
    def wrapped_grad(argnum, g, ans, vs, gvs, args, kwargs):
        return make_vjp(fun, argnum)(*args, **kwargs)[0](g)
    wrapped = primitive(fun)
    wrapped.vjp = wrapped_grad
    return wrapped