code: def test_jax_vmap(
      func,
      dtype_and_arrays_and_axes,
      in_axes_as_cont,
      backend_fw,
  ):
      dtype, generated_arrays, in_axes = dtype_and_arrays_and_axes
      ivy.set_backend(backend_fw)
      arrays = [ivy.native_array(array) for array in generated_arrays]
      if in_axes_as_cont:
          vmapped_func = vmap(func, in_axes=in_axes, out_axes=0)
      else:
          vmapped_func = vmap(func, in_axes=0, out_axes=0)
  
      assert callable(vmapped_func)
  
      try:
          fw_res = helpers.flatten_and_to_np(
              ret=vmapped_func(*arrays), backend=backend_fw
          )
          fw_res = fw_res if len(fw_res) else None
      except Exception:
          fw_res = None
  
      ivy.previous_backend()
  
      ivy.set_backend("jax")
      arrays = [ivy.native_array(array) for array in generated_arrays]
      if in_axes_as_cont:
          jax_vmapped_func = jax.vmap(func, in_axes=in_axes, out_axes=0)
      else:
          jax_vmapped_func = jax.vmap(func, in_axes=0, out_axes=0)
  
      assert callable(jax_vmapped_func)
  
      try:
          jax_res = helpers.flatten_and_to_np(
              ret=jax_vmapped_func(*arrays), backend="jax"
          )
          jax_res = jax_res if len(jax_res) else None
      except Exception:
          jax_res = None
  
      ivy.previous_backend()
  
      if fw_res is not None and jax_res is not None:
          helpers.value_test(
              ret_np_flat=fw_res,
              ret_np_from_gt_flat=jax_res,
              rtol=1e-1,
              atol=1e-1,
              backend=backend_fw,
              ground_truth_backend="jax",
          )
  
      elif fw_res is None and jax_res is None:
          pass
      else:
          assert False, "One of the results is None while other isn't"
  
file_path: ivy/ivy_tests/test_ivy/test_frontends/test_jax/test_general_functions.py
imports: import ivy
  import ivy_tests.test_ivy.helpers
  from ivy_tests.test_ivy.helpers import handle_frontend_test, BackendHandler
  from ivy.functional.frontends.jax import vmap
  from hypothesis import strategies
  import jax
label_desc: Use image node_id 6 for calling a global function with example usage: test_jax_vmap(func, dtype_and_arrays_and_axes, in_axes_as_cont, backend_fw) without return types