code: def _fn_with_custom_grad(
      fn,
      inputs,
      grad_fn,
      use_global_vars=False,
      use_entire_scope=False,
      train_vars=None,
  ):
      """Create a subgraph with a custom gradient.
      Args:
        fn: function that takes inputs as arguments and produces 1 or more Tensors.
        inputs: list<Tensor>, will be passed as fn(*inputs).
        grad_fn: function with signature
          (inputs, vars, 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:
        fn(*inputs)
      """
      vs = tf.compat.v1.get_variable_scope()
  
      # Use a function attribute to store the local variables. This means that on graph rebuild,
      # When variables already exist we know what variables relate to this function.
      get_vars_fn = (
          vs.global_variables
          if use_global_vars
          else vs.trainable_variables
      )
      len_before_vars = len(get_vars_fn())
      inputs = list(inputs)
      outputs = fn(*inputs)
      if not train_vars:
          if use_entire_scope:
              train_vars = get_vars_fn()
          else:
              train_vars = get_vars_fn()[len_before_vars:]
  
      if grad_fn is None:
          return outputs
  
      if not isinstance(outputs, (tuple, list)):
          outputs = [outputs]
      outputs = list(outputs)
      raw_outputs = outputs
      outputs = [o for o in outputs if isinstance(o, tf.Tensor)]
  
      defun_inputs = [inputs, train_vars, outputs]
  
      def custom_grad_fn(op, *dys):
          """Custom grad fn applying grad_fn for identity Defun."""
          fn_inputs, fn_vars, fn_outputs = tf.nest.pack_sequence_as(
              defun_inputs, list(op.inputs)
          )
          dys = list(dys)
          assert len(fn_outputs) == len(outputs)
          assert len(fn_outputs) == len(dys)
  
          grad_inputs, grad_vars = grad_fn(
              fn_inputs, fn_vars, fn_outputs, dys
          )
          grad_outputs = [None] * len(fn_outputs)
          return tuple(grad_inputs + grad_vars + grad_outputs)
  
      # The Defun takes as input the original inputs, the trainable variables
      # created in fn, and the outputs. In the forward it passes through the
      # outputs. In the backwards, it produces gradients for the original inputs
      # and the trainable variables.
      in_types = [t.dtype for t in inputs]
      out_types = [t.dtype for t in outputs]
      var_types = [t.dtype for t in train_vars]
  
      @function.Defun(
          *(in_types + var_types + out_types),
          func_name="identity_custom_grad%d" % ops.uid(),
          python_grad_func=custom_grad_fn,
          shape_func=lambda _: [t.get_shape() for t in outputs],
      )
      def identity(*args):
          _, _, outs = tf.nest.pack_sequence_as(defun_inputs, args)
          return tuple(
              [
                  tf.identity(t) if isinstance(t, tf.Tensor) else t
                  for t in outs
              ]
          )
  
      flat_inputs = tf.nest.flatten(defun_inputs)
      id_out = identity(*flat_inputs)
      if isinstance(id_out, tuple):
          id_out = list(id_out)
      else:
          id_out = [id_out]
      combined_out = []
      for o in raw_outputs:
          if isinstance(o, tf.Tensor):
              combined_out.append(id_out.pop(0))
          else:
              combined_out.append(o)
      assert len(id_out) == 0, str(id_out)
      return tuple(combined_out)
  
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 2 for calling a global function with example usage: _fn_with_custom_grad(fn, inputs, grad_fn, use_global_vars, use_entire_scope, train_vars) and returns: tuple, outputs, tuple, tuple