code: def test_map_non_jax():
      # this contains a non-jax value for the `activation` field
      # and will therefore break filter_vmap if not filtered out
      # at input and output
      pytree = eqx.nn.MLP(
          2,
          2,
          2,
          2,
          activation=jax.nn.relu,
          key=jax.random.PRNGKey(42),
      )
  
      def identity(x):
          """will return a pytree with non-jax fields, which could break filter_vmap"""
          return x
  
      _ = eqx.filter_vmap(identity)(pytree)
  
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 9 for calling a global function with example usage: test_map_non_jax() and returns: x