code: def test_default():
      @filter_pmap(in_axes=_zero_if_inexact_array_else_none)
      def f(a, b):
          with jax.numpy_dtype_promotion("standard"):
              return a + b
  
      assert shaped_allclose(
          f(jnp.array(3), jnp.array([3.0])), jnp.array([6.0])
      )
  
      with pytest.raises(ValueError):
          f(jnp.array(3.0), jnp.array([3.0]))
  
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 4 for calling a global function with example usage: test_default() and returns: unknown