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.
          """
          wrapped_function = _autograd.extend.primitive(func)
  
          def wrapped_grad_func(i, ans, *args, **kwargs):
              grads = grad_funcs[i](*args, **kwargs)
              if isinstance(grads, float):
                  return lambda g: g * grads
              if grads.ndim == 2:
                  return lambda g: g[..., None] * grads
              if grads.ndim == 3:
                  return lambda g: g[..., None, None] * grads
              return lambda g: g * grads
  
          if len(grad_funcs) == 1:
              _autograd.extend.defvjp(
                  wrapped_function,
                  lambda ans, *args, **kwargs: wrapped_grad_func(
                      0, ans, *args, **kwargs
                  ),
              )
          elif len(grad_funcs) == 2:
              _autograd.extend.defvjp(
                  wrapped_function,
                  lambda ans, *args, **kwargs: wrapped_grad_func(
                      0, ans, *args, **kwargs
                  ),
                  lambda ans, *args, **kwargs: wrapped_grad_func(
                      1, ans, *args, **kwargs
                  ),
              )
          elif len(grad_funcs) == 3:
              _autograd.extend.defvjp(
                  wrapped_function,
                  lambda ans, *args, **kwargs: wrapped_grad_func(
                      0, ans, *args, **kwargs
                  ),
                  lambda ans, *args, **kwargs: wrapped_grad_func(
                      1, ans, *args, **kwargs
                  ),
                  lambda ans, *args, **kwargs: wrapped_grad_func(
                      2, ans, *args, **kwargs
                  ),
              )
          else:
              raise NotImplementedError(
                  "custom_gradient is not yet implemented "
                  "for more than 3 gradients."
              )
  
          return wrapped_function
  
      return decorator
  
file_path: geomstats/geomstats/_backend/autograd/autodiff.py
imports: import autograd
  import autograd.numpy
  from autograd import jacobian
label_desc: Use image node_id 1 for calling a global function with example usage: custom_gradient() and returns: decorator, wrapped_function, function, function, function, function