code: def test_vmap(getkey):
      mlp = eqx.nn.MLP(2, 2, 512, 2, key=getkey())
      mlp_noinline = eqx.internal.noinline(mlp)
      mlp_vmap = jax.vmap(mlp)
      mlp_jit_vmap = eqx.filter_jit(mlp_vmap, donate="none")
      mlp_vmap_noinline = jax.vmap(mlp_noinline)
      mlp_jit_vmap_noinline = eqx.filter_jit(
          mlp_vmap_noinline, donate="none"
      )
      x = jr.normal(getkey(), (5, 2))
      o1 = mlp_vmap(x)
      o2 = mlp_jit_vmap(x)
      o3 = mlp_vmap_noinline(x)
      o4 = mlp_jit_vmap_noinline(x)
      assert tree_allclose(o1, o2, atol=1e-5)
      assert tree_allclose(o1, o3, atol=1e-5)
      assert tree_allclose(o1, o4, atol=1e-5)
  
file_path: equinox/tests/test_noinline.py
imports: import equinox
  import jax
  import jax.numpy
  import jax.random
  from .helpers import tree_allclose
label_desc: Use image node_id 3 for calling a global function with example usage: test_vmap(getkey) without return types