code: def _grad(func, argnums=0):
      def _wrapped_grad(*x, **kwargs):
          if not hasattr(x[0], "ndim") or x[0].ndim < 2:
              return _autograd.grad(func, argnum=argnums)(*x, **kwargs)
  
          return _autograd.elementwise_grad(func, argnum=argnums)(
              *x, **kwargs
          )
  
      return _wrapped_grad
  
file_path: geomstats/geomstats/_backend/autograd/autodiff.py
imports: import autograd
  import autograd.numpy
  from autograd import jacobian
label_desc: Use image node_id 2 for calling a global function with example usage: _grad(func, argnums) and returns: _wrapped_grad, _autograd, _autograd