code: def _recompute_grad(fn, args):
    
  
    cached_vs = []
    cached_arg_scope = []
  
    def grad_fn(inputs, variables, outputs, output_grads):
      """Recompute outputs for gradient computation."""
      del outputs
      variables = [underlying_variable_ref(v) for v in variables]
      # Recompute outputs
      with tf.control_dependencies(output_grads):
        with tf.contrib.framework.arg_scope(cached_arg_scope[0]):
          with tf.variable_scope(cached_vs[0], reuse=True):
            outputs = fn(*inputs)
  
      if not isinstance(outputs, (list, tuple)):
        outputs = [outputs]
      outputs = list(outputs)
      grads = tf.gradients(outputs, inputs + variables, output_grads)
      grad_inputs = grads[:len(inputs)]
      grad_vars = grads[len(inputs):]
      # TODO(rsepassi): Make fn_with_custom_grad work with bfloat16.
      # If the input gradients are bfloat16, it's assumed the variables are
      # bfloat16. This is a hack to ensure that grad_vars are the right type.
      if grad_inputs[0].dtype == tf.bfloat16:
        grad_vars = [tf.cast(grad_var, tf.bfloat16) for grad_var in grad_vars]
      return grad_inputs, grad_vars
  
    @fn_with_custom_grad(grad_fn)
    def fn_with_recompute(*args):
      cached_vs.append(tf.get_variable_scope())
      cached_arg_scope.append(tf.contrib.framework.current_arg_scope())
      return fn(*args)
  
    return fn_with_recompute(*args)
docs: See recompute_grad.