code: def _recompute_grad(f):
      """An eager-compatible version of recompute_grad.
  
      For f(*args, **kwargs), this supports gradients with respect to args or
      kwargs, but kwargs are currently only supported in eager-mode.
      Note that for keras layer and model objects, this is handled automatically.
  
      Warning: If `f` was originally a tf.keras Model or Layer object, `g` will not
      be able to access the member variables of that object, because `g` returns
      through the wrapper function `inner`.  When recomputing gradients through
      objects that inherit from keras, we suggest keeping a reference to the
      underlying object around for the purpose of accessing these variables.
  
      Args:
        f: function `f(*x)` that returns a `Tensor` or sequence of `Tensor` outputs.
  
      Returns:
       A function `g` that wraps `f`, but which recomputes `f` on the backwards
       pass of a gradient call.
      """
  
      @tf1.custom_gradient
      def inner(*args, **kwargs):
          """Inner function closure for calculating gradients."""
          current_var_scope = tf1.get_variable_scope()
          with tape_lib.stop_recording():
              result = f(*args, **kwargs)
  
          def grad_wrapper(*wrapper_args, **grad_kwargs):
              """Wrapper function to accomodate lack of kwargs in graph mode decorator."""
  
              @tf1.custom_gradient
              def inner_recompute_grad(*dresult):
                  """Nested custom gradient function for computing grads in reverse and forward mode autodiff."""
                  # Gradient calculation for reverse mode autodiff.
                  variables = grad_kwargs.get("variables")
                  with tf.GradientTape() as t:
                      id_args = tf.nest.map_structure(tf.identity, args)
                      t.watch(id_args)
                      if variables is not None:
                          t.watch(variables)
                      with tf.control_dependencies(dresult):
                          with tf1.variable_scope(current_var_scope):
                              result = f(*id_args, **kwargs)
                  kw_vars = []
                  if variables is not None:
                      kw_vars = list(variables)
                  grads = t.gradient(
                      result,
                      list(id_args) + kw_vars,
                      output_gradients=dresult,
                      unconnected_gradients=tf.UnconnectedGradients.ZERO,
                  )
  
                  def transpose(*t_args, **t_kwargs):
                      """Gradient function calculation for forward mode autodiff."""
                      # Just throw an error since gradients / activations are not stored on
                      # tape for recompute.
                      raise NotImplementedError(
                          "recompute_grad tried to transpose grad of {}. "
                          "Consider not using recompute_grad in forward mode"
                          "autodiff".format(f.__name__)
                      )
  
                  return (
                      grads[: len(id_args)],
                      grads[len(id_args) :],
                  ), transpose
  
              return inner_recompute_grad(*wrapper_args)
  
          return result, grad_wrapper
  
      return inner
  
imports: from typing import Text, Tuple, Union
  import tensorflow
  import argparse
  import tensorflow.compat.v1
  from enum import Enum
  from collections import namedtuple
  from tensorflow.python.eager import tape
  from tensorflow.python.tpu import tpu_function
label_desc: Use image node_id 11 for calling a global function with example usage: _recompute_grad(f) and returns: inner, result,  grad_wrapper, inner_recompute_grad, grads,  grads,  transpose