code: def _recompute_grad(
      fn, args, use_data_dep=_USE_DEFAULT, tupleize_grads=False
  ):
      """See recompute_grad."""
      has_is_recompute_kwarg = (
          "is_recomputing" in tf_inspect.getargspec(fn).args
      )
      for arg in args:
          if not isinstance(arg, framework_ops.Tensor):
              raise ValueError("All inputs to function must be Tensors")
      use_data_dep_ = use_data_dep
      if use_data_dep_ == _USE_DEFAULT:
          use_data_dep_ = _is_on_tpu()
  
      # Use custom_gradient and return a grad_fn that recomputes on the backwards
      # pass.
      @custom_gradient.custom_gradient
      def fn_with_recompute(*args):
          """Wrapper for fn."""
          # Capture the variable and arg scopes so we can re-enter them when
          # recomputing.
          vs = variable_scope.get_variable_scope()
          arg_scope = contrib_framework_ops.current_arg_scope()
          # Track all variables touched in the function.
          with backprop.GradientTape() as tape:
              fn_kwargs = {}
              if has_is_recompute_kwarg:
                  fn_kwargs["is_recomputing"] = False
              outputs = fn(*args, **fn_kwargs)
          original_vars = set(tape.watched_variables())
  
          def _grad_fn(output_grads, variables=None):
              # Validate that custom_gradient passes the right variables into grad_fn.
              if original_vars:
                  assert variables, (
                      "Fn created variables but the variables were not "
                      "passed to the gradient fn."
                  )
                  if set(variables) != original_vars:
                      raise ValueError(_WRONG_VARS_ERR)
  
              return _recomputing_grad_fn(
                  compute_fn=fn,
                  original_args=args,
                  original_vars=original_vars,
                  output_grads=output_grads,
                  grad_fn_variables=variables,
                  use_data_dep=use_data_dep_,
                  tupleize_grads=tupleize_grads,
                  arg_scope=arg_scope,
                  var_scope=vs,
                  has_is_recompute_kwarg=has_is_recompute_kwarg,
              )
  
          # custom_gradient inspects the signature of the function to determine
          # whether the user expects variables passed in the grad_fn. If the function
          # created variables, the grad_fn should accept the "variables" kwarg.
          if original_vars:
  
              def grad_fn(*output_grads, **kwargs):
                  return _grad_fn(output_grads, kwargs["variables"])
  
          else:
  
              def grad_fn(*output_grads):
                  return _grad_fn(output_grads)
  
          return outputs, grad_fn
  
      return fn_with_recompute(*args)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import re
  import numpy
  import six
  from six.moves import xrange
  from ..None import contrib_framework
  from tensorflow.python.eager import backprop
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.layers import base
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.ops import custom_gradient
  from tensorflow.python.ops import gradients_impl
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.util import nest
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 12 for calling a global function with example usage: _recompute_grad(fn, args, use_data_dep, tupleize_grads) and returns: fn_with_recompute, outputs,  grad_fn, _recomputing_grad_fn, _grad_fn, _grad_fn