code: def test_num_traces():
      num_traces = 0
  
      @filter_pmap(in_axes=None, axis_size=1)
      def f(x, y):
          nonlocal num_traces
          num_traces += 1
          return x + y
  
      assert shaped_allclose(f(jnp.array(1), 2), jnp.array([3]))
      assert num_traces == 2  # eval_shape + pmap
      assert shaped_allclose(f(jnp.array(3), 2), jnp.array([5]))
      assert num_traces == 2
      assert shaped_allclose(f(jnp.array(1), 3), jnp.array([4]))
      assert num_traces == 4
      assert shaped_allclose(f(jnp.array([3]), 3), jnp.array([[6]]))
      assert num_traces == 6
      assert shaped_allclose(f(jnp.array([4]), 3), jnp.array([[7]]))
      assert num_traces == 6
  
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 7 for calling a global function with example usage: test_num_traces() and returns: unknown