code: def test_out():
      def f(x):
          return x
  
      o1 = eqx.filter_vmap(f, out_axes=None, axis_size=5)(1)
      o2 = eqx.filter_vmap(f, in_axes=None, out_axes=None, axis_size=5)(
          jnp.array([3, 4])
      )
      o3 = eqx.filter_vmap(f, out_axes=0, axis_size=5)(1)
      o4 = eqx.filter_vmap(f, in_axes=None, out_axes=0, axis_size=5)(
          jnp.array([3, 4])
      )
  
      assert tree_allclose(o1, 1)
      assert tree_allclose(o2, jnp.array([3, 4]))
      assert tree_allclose(o3, jnp.array([1, 1, 1, 1, 1]))
      assert tree_allclose(
          o4, jnp.array([[3, 4], [3, 4], [3, 4], [3, 4], [3, 4]])
      )
  
file_path: equinox/tests/test_vmap.py
imports: from typing import Union
  import equinox
  import jax
  import jax.numpy
  import jax.random
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 4 for calling a global function with example usage: test_out() and returns: x