code: def test_vmap():
      vf = jax.vmap(_f)
      jvf = jax.jit(vf)
      good = jnp.array([1.0, 1.0])
      bad1 = jnp.array([1.0, -1.0])
      bad2 = jnp.array([-1.0, -1.0])
  
      vf(good)
      jvf(good)
      with _error:
          vf(bad1)
      with _error:
          vf(bad2)
      with _error:
          jvf(bad1)
      with _error:
          jvf(bad2)
  
file_path: equinox/tests/test_errors.py
imports: import equinox
  import equinox.internal
  import jax
  import jax.numpy
  import pytest
label_desc: Use image node_id 3 for calling a global function with example usage: test_vmap() without return types