code: def test_named_reduction():
      def f(x):
          y = x + 1
          return jax.lax.psum(y, axis_name="device")
  
      n = jax.local_device_count()
      output = filter_pmap(f, axis_name="device")(jnp.zeros(n))
  
      assert shaped_allclose(output, n * jnp.ones(n))
  
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 11 for calling a global function with example usage: test_named_reduction() and returns: jax