code: def test_args():
      @filter_pmap(
          in_axes=(
              _zero_if_inexact_array_else_none,
              [{"a": None}],
              0,
              eqx.if_array(0),
          )
      )
      def f(a, b, c, d):
          return a + b[0]["a"] + c + d
  
      out = f(
          jnp.array([1]), [{"a": jnp.array([2])}], jnp.array([3]), 4
      )
      assert shaped_allclose(out, jnp.array([[10]]))
  
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 3 for calling a global function with example usage: test_args() and returns: unknown