code: def test_max_traces_clone(getkey):
      lin = eqx.nn.Linear(3, 4, key=getkey())
      lin = eqx.filter_jit(
          eqx.debug.assert_max_traces(lin, max_traces=2)
      )
      leaves, treedef = jtu.tree_flatten(lin)
      lin2 = jtu.tree_unflatten(treedef, leaves)
  
      lin(jnp.array([1.0, 2.0, 3.0]))
      assert eqx.debug.get_num_traces(lin) == 1
      assert eqx.debug.get_num_traces(lin2) == 1
      with jax.numpy_dtype_promotion("standard"):
          lin2(jnp.array([1, 2, 3]))
          assert eqx.debug.get_num_traces(lin) == 2
          assert eqx.debug.get_num_traces(lin2) == 2
  
          with pytest.raises(
              RuntimeError, match="can only be traced 2 times"
          ):
              lin(jnp.array([False, False, False]))
  
file_path: equinox/tests/test_debug.py
imports: import equinox
  import jax
  import jax.numpy
  import jax.tree_util
  import pytest
label_desc: Use image node_id 4 for calling a global function with example usage: test_max_traces_clone(getkey) without return types