code: def checkpoint(fun):
      """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, ans, args, kwargs):
          return make_vjp(fun, argnum)(*args, **kwargs)[0]
  
      wrapped = primitive(fun)
      defvjp_argnum(wrapped, wrapped_grad)
      return wrapped
  
imports: from __future__ import absolute_import
  from functools import partial
  from collections import OrderedDict
  import warnings
  from .wrap_util import unary_to_nary
  from .builtins import tuple
  from .core import make_vjp, make_jvp
  from .extend import primitive, defvjp_argnum, vspace
  import autograd.numpy
label_desc: Use image node_id 16 for calling a global function with example usage: checkpoint(fun) and returns: wrapped, make_vjp