code: def fn_with_custom_grad(
      grad_fn,
      use_global_vars=False,
      use_entire_scope=False,
      train_vars=None,
  ):
      """Decorator to create a subgraph with a custom gradient function.
      The subgraph created by the decorated function is NOT put in a Defun and so
      does not suffer from the limitations of the Defun (all subgraph ops on the
      same device, no summaries).
      Args:
        grad_fn: function with signature
          (inputs, variables, outputs, output_grads) -> (grad_inputs, grad_vars),
          all of which are lists of Tensors.
        use_global_vars: if True, variables will be the global variables created.
          If False, will be the trainable variables.
      Returns:
        Decorator for function such that the gradient is defined by grad_fn.
      """
  
      def dec(fn):
          @functools.wraps(fn)
          def wrapped(*args):
              return _fn_with_custom_grad(
                  fn,
                  args,
                  grad_fn,
                  use_global_vars=use_global_vars,
                  use_entire_scope=use_entire_scope,
                  train_vars=None,
              )
  
          return wrapped
  
      return dec
  
file_path: finetune/finetune/optimizers/recompute_grads.py
imports: import functools
  import tensorflow
  from tensorflow.python.framework import function, ops
  from finetune.util.optimize_loss import is_numerical_tensor
label_desc: Use image node_id 1 for calling a global function with example usage: fn_with_custom_grad(grad_fn, use_global_vars, use_entire_scope, train_vars) and returns: dec, wrapped, _fn_with_custom_grad