code: def test_keyword_in_axes(getkey):
      x = jr.normal(getkey(), (3, 4))
      y = jr.normal(getkey(), (1, 3))
      out = eqx.filter_vmap(lambda x, y: x + y, in_axes=dict(y=1))(x, y)
      true_out = x + y.T
      assert tree_allclose(out, true_out)
  
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 10 for calling a global function with example usage: test_keyword_in_axes(getkey) without return types