code: def test_ensemble(getkey):
      def make(key):
          return eqx.nn.MLP(5, 4, 3, 2, key=getkey())
  
      keys = jr.split(getkey(), 7)
      models = eqx.filter_vmap(make)(keys)
  
      def call(model, x):
          return model(x)
  
      xs1 = jr.normal(getkey(), (7, 5))
      assert eqx.filter_vmap(call)(models, xs1).shape == (7, 4)
  
      xs2 = jr.normal(getkey(), (5,))
      assert eqx.filter_vmap(call, in_axes=(eqx.if_array(0), None))(
          models, xs2
      ).shape == (7, 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 6 for calling a global function with example usage: test_ensemble(getkey) and returns: eqx, model