code: def test_map_non_jax():
      devices = jax.local_devices()
  
      # this contains a non-jax value for the `activation` field
      # and will therefore break filter_pmap 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 maybe_replicate(value):
          if eqx.is_array(value):
              return jax.device_put_replicated(value, devices)
          else:
              return value
  
      pytree_sharded = jtu.tree_map(maybe_replicate, pytree)
  
      def identity(x):
          """will return a pytree with non-jax fields, which could break filter_pmap"""
          return x
  
      filter_pmap(identity)(pytree_sharded)
  
imports: import functools
  from typing import Any, Union
  import equinox
  import jax
  import jax.numpy
  import jax.random
  import jax.tree_util
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 12 for calling a global function with example usage: test_map_non_jax() and returns: x, jax, value