code: def custom_gradient(*grad_funcs):
      """Create a decorator that allows a function to define its custom gradient(s).
  
      Parameters
      ----------
      *grad_funcs : callables
          Custom gradient functions.
  
      Returns
      -------
      decorator : callable
          This decorator, used on any function func, associates the
          input grad_funcs as the gradients of func.
      """
  
      def decorator(func):
          """Decorate a function to define its custome gradient(s).
  
          Parameters
          ----------
          func : callable
              Function whose gradients will be assigned by grad_funcs.
  
          Returns
          -------
          wrapped_function : callable
              Function func with gradients specified by grad_funcs.
          """
  
          class func_with_grad(_torch.autograd.Function):
              """Wrapper class for a function with custom grad."""
  
              @staticmethod
              def forward(ctx, *args):
                  ctx.save_for_backward(*args)
                  return func(*args)
  
              @staticmethod
              def backward(ctx, grad_output):
                  inputs = ctx.saved_tensors
  
                  grads = ()
                  for custom_grad in grad_funcs:
                      grads = (
                          *grads,
                          grad_output * custom_grad(*inputs),
                      )
  
                  if len(grads) == 1:
                      return grads[0]
                  return grads
  
          def wrapped_function(*args, **kwargs):
              new_inputs = args + tuple(kwargs.values())
              return func_with_grad.apply(*new_inputs)
  
          return wrapped_function
  
      return decorator
  
file_path: geomstats/geomstats/_backend/pytorch/autodiff.py
imports: import functools
  import numpy
  import torch
  from torch.autograd.functional import hessian
  from torch.autograd.functional import jacobian
label_desc: Use image node_id 1 for calling a global function with example usage: custom_gradient() and returns: decorator, wrapped_function, func_with_grad, func, grads, grads