code: def test_out():
      def f(x):
          return x
  
      o1 = filter_pmap(f, in_axes=None, out_axes=None, axis_size=1)(
          jnp.array([3, 4])
      )
      o2 = filter_pmap(f, out_axes=0, axis_size=1)(1)
      o3 = filter_pmap(f, in_axes=None, out_axes=0, axis_size=1)(
          jnp.array([3, 4])
      )
  
      assert shaped_allclose(o1, jnp.array([3, 4]))
      assert shaped_allclose(o2, jnp.array([1]))
      assert shaped_allclose(o3, jnp.array([[3, 4]]))
  
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 5 for calling a global function with example usage: test_out() and returns: x