docstr: Test that jacobian_vec is correctly vectorized.
  
  The autodiff jacobian is not vectorized by default in torch, tf and autograd.
  
  The jacobian of a function f going from an input space A to an output
  space B is a matrix of shape (dim_B, dim_A).
  - The columns index the derivatives wrt. the coordinates of the input space A.
  - The rows index the coordinates of the output space B.
docstr_tok: ["Test","that","jacobian_vec","is","correctly","vectorized",".","The","autodiff","jacobian","is","not","vectorized","by","default","in","torch",",","tf","and","autograd",".","The","jacobian","of","a","function","f","going","from","an","input","space","A","to","an","output","space","B","is","a","matrix","of","shape","(","dim_B",",","dim_A",")",".","-","The","columns","index","the","derivatives","wrt",".","the","coordinates","of","the","input","space","A",".","-","The","rows","index","the","coordinates","of","the","output","space","B","."]
code: def test_jacobian_vec(self):
      """Test that jacobian_vec is correctly vectorized.
  
      The autodiff jacobian is not vectorized by default in torch, tf and autograd.
  
      The jacobian of a function f going from an input space A to an output
      space B is a matrix of shape (dim_B, dim_A).
      - The columns index the derivatives wrt. the coordinates of the input space A.
      - The rows index the coordinates of the output space B.
      """
      radius = 4.0
      embedding_dim, dim = 3, 2
  
      points = gs.array([[gs.pi / 3, gs.pi], [gs.pi / 5, gs.pi / 2]])
      thetas = points[:, 0]
      phis = points[:, 1]
      jacobian_ai = gs.autodiff.jacobian_vec(_sphere_immersion)(points)
  
      expected_1i = gs.stack(
          [
              gs.array(
                  [
                      radius * gs.cos(theta) * gs.cos(phi),
                      -radius * gs.sin(theta) * gs.sin(phi),
                  ]
              )
              for theta, phi in zip(thetas, phis)
          ]
      )
      expected_2i = gs.stack(
          [
              gs.array(
                  [
                      radius * gs.cos(theta) * gs.sin(phi),
                      radius * gs.sin(theta) * gs.cos(phi),
                  ]
              )
              for theta, phi in zip(thetas, phis)
          ]
      )
      expected_3i = gs.stack(
          [
              gs.array(
                  [
                      -radius * gs.sin(theta),
                      0,
                  ]
              )
              for theta in thetas
          ]
      )
      expected_ai = gs.stack(
          [expected_1i, expected_2i, expected_3i], axis=1
      )
      self.assertAllClose(
          jacobian_ai.shape, (len(points), embedding_dim, dim)
      )
      self.assertAllClose(jacobian_ai.shape, expected_ai.shape)
      self.assertAllClose(jacobian_ai, expected_ai)
  
imports: import warnings
  import numpy
  import pytest
  import geomstats.backend
  from geomstats.geometry.grassmannian import Grassmannian
  from geomstats.geometry.special_euclidean import SpecialEuclidean
  from geomstats.test.test_case import TestCase, autograd_and_torch_only, autograd_only, np_only
label_desc: Use image node_id 20 for calling the TestAutodiff obj's underlying member method code with example usage: obj.test_jacobian_vec() without return types