code: def jacobian_vec(func, point_ndim=1):
      """Return a function that returns the jacobian of func.
  
      We note that the jacobian function of torch is not vectorized
      by default, thus we modify its behavior here.
  
      Default pytorch behavior:
  
      If the jacobian for one point of shape (2,) is of shape (3, 2),
      then calling the jacobian on 4 points with shape (4, 2) will
      be of shape (3, 2, 4, 2).
  
      Modified behavior:
  
      Calling the jacobian on 4 points gives a tensor of shape (4, 3, 2).
  
      We use a for-loop to allow this function to be vectorized with
      respect to several inputs in point, because the flag vectorize=True
      fails.
  
      Parameters
      ----------
      func : callable
          Function whose jacobian is computed.
  
      Returns
      -------
      _ : callable
          Function taking point as input and returning
          the jacobian of func at point.
      """
  
      def _jacobian(point):
          if point.ndim == point_ndim:
              return _torch_jacobian(
                  func=lambda x: func(x), inputs=point
              )
          return _torch.stack(
              [
                  _torch_jacobian(
                      func=lambda x: func(x), inputs=one_point
                  )
                  for one_point in point
              ],
              axis=0,
          )
  
      return _jacobian
  
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 3 for calling a global function with example usage: jacobian_vec(func, point_ndim) and returns: _jacobian, _torch, _torch_jacobian