code: def test_grad_jit_new():
      num_traces = 0
  
      @eqx.filter_custom_vjp
      def f(x):
          return x
  
      @f.def_fwd
      def f_fwd(perturbed, x):
          del perturbed
          return x, None
  
      @f.def_bwd
      def f_bwd(residual, g, perturbed, x):
          del residual, perturbed, x
          nonlocal num_traces
          num_traces += 1
          return g + 2
  
      x = jnp.array(1.0)
  
      jitf = jax.jit(f)
      assert eqx.filter_grad(jitf)(x) == 3
      assert eqx.filter_grad(jitf)(x) == 3
      assert num_traces == 1
      assert (
          eqx.filter_grad(eqx.filter_jit(f))(x) == 3
      )  # pyright: ignore
      assert (
          eqx.filter_grad(eqx.filter_jit(f))(x) == 3
      )  # pyright: ignore
      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 numpy
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 5 for calling a global function with example usage: test_grad_jit_new() and returns: x, x,  None, unknown