code: def assert_max_traces(
      fn: Callable = sentinel, *, max_traces: Optional[int]
  ):
      """Asserts that the wrapped callable is not called more than `max_traces` times.
  
      The typical use-case for this is to check that a JIT-compiled function is not
      compiled more than `max_traces` times. (I.e. this function can be used to guard
      against bugs.) In this case it should be placed within the JIT wrapper.
  
      !!! Example
  
          ```python
          @eqx.filter_jit
          @eqx.debug.assert_max_traces(max_traces=1)
          def f(x, y):
              return x + y
          ```
  
      **Arguments:**
  
      - `fn`: The callable to wrap.
      - `max_traces`: keyword only argument. The maximum number of calls that are
          allowed. Can be `None` to allow arbitrarily many calls; in this case the number
          of calls can still can be found via [`equinox.debug.get_num_traces`][].
  
      **Returns:**
  
      A wrapped version of `fn` that tracks the number of times it is called. This will
      raise a `RuntimeError` if is called more than `max_traces` many times.
  
      !!! info
  
          See also
          [`chex.assert_max_traces`](https://github.com/google-deepmind/chex#jax-tracing-assertions)
          which provides similar functionality.
  
          The differences are that (a) Chex's implementation is a bit stricter, as the
          following will raise:
          ```python
          import chex
          import jax
  
          def f(x):
              pass
  
          f2 = jax.jit(chex.assert_max_traces(f, 1))
          f3 = jax.jit(chex.assert_max_traces(f, 1))
  
          f2(1)
          f3(1)  # will raise, despite the fact that f2 and f3 are different.
          ```
          and (b) Equinox's implementation supports non-function callables.
  
          You may prefer the Chex version if you prefer the stricter raising behaviour of
          the above code.
      """
      if fn is sentinel:
          return ft.partial(
              assert_max_traces, max_traces=max_traces
          )  # pyright: ignore
      return module_update_wrapper(_AssertMaxTraces(fn, max_traces))
  
imports: import functools
  import inspect
  import weakref
  from collections.abc import Callable
  from typing import Optional, overload, TypeVar
  from typing_extensions import ParamSpec
  from .._custom_types import sentinel
  from .._eval_shape import filter_eval_shape
  from .._module import field, Module, module_update_wrapper
  from .._pretty_print import tree_pformat
  from .._tree import tree_equal
label_desc: Use image node_id 3 for calling a global function with example usage: assert_max_traces(fn) and returns: module_update_wrapper, ft