code: def test_num_traces():
      num_traces = 0
  
      def f(x):
          nonlocal num_traces
          num_traces += 1
  
      f = eqx.filter_jit(f)
  
      f(jnp.zeros(2))
      f(jnp.zeros(2))
      assert num_traces == 1
  
      f(jnp.zeros(3))
      f(jnp.zeros(3))
      assert num_traces == 2
  
      f([jnp.zeros(2)])
      f([jnp.zeros(2), jnp.zeros(3)])
      f([jnp.zeros(2), True])
      assert num_traces == 5
  
      num_traces = 0
  
      def g(x, y):
          nonlocal num_traces
          num_traces += 1
  
      g = eqx.filter_jit(g)
  
      g(jnp.zeros(2), True)
      g(jnp.zeros(2), False)
      assert num_traces == 2
  
      num_traces = 0
  
      def h(x, y, z, w):
          nonlocal num_traces
          num_traces += 1
  
      h = eqx.filter_jit(h)
  
      h(True, {"a": 1, "b": 1}, True, True)
      h(False, {"a": 1, "b": 1}, True, True)
      h(True, {"a": 1, "b": 0}, True, True)
      h(True, {"a": 1, "b": 1}, True, 2)
      h(True, {"a": 1, "b": 1}, 5, True)
      assert num_traces == 5
      h(True, {"a": 1, "b": 1}, True, True)
      assert num_traces == 5
      h(True, {"a": 2, "b": 1}, True, True)
      assert num_traces == 6
  
imports: import warnings
  from typing import 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 2 for calling a global function with example usage: test_num_traces() without return types