code: def test_jit_grad():
      num_traces = 0
  
      def f(x):
          nonlocal num_traces
          num_traces += 1
          return x + 1
  
      assert eqx.filter_jit(eqx.filter_grad(f))(jnp.array(1.0)) == 1
      assert eqx.filter_jit(eqx.filter_grad(f))(jnp.array(2.0)) == 1
      assert num_traces == 1
  
      assert eqx.filter_jit(eqx.filter_value_and_grad(f))(
          jnp.array(1.0)
      ) == (2, 1)
      assert eqx.filter_jit(eqx.filter_value_and_grad(f))(
          jnp.array(2.0)
      ) == (3, 1)
      assert num_traces == 2
  
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 6 for calling a global function with example usage: test_jit_grad() and returns: unknown