code: def fn_with_custom_grad(grad_fn, use_global_vars=False):
    
  
    def dec(fn):
  
      @functools.wraps(fn)
      def wrapped(*args):
        return _fn_with_custom_grad(
            fn, args, grad_fn, use_global_vars=use_global_vars)
  
      return wrapped
  
    return dec
docs: 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.