code: def test_filter_jvp():
      _map_is_array_like = lambda pytree: jax.tree_map(
          lambda x: jnp.array(x) if eqx.is_array_like(x) else x, pytree
      )
  
      def after_filter_jit(fun):
          fun_jit = eqx.filter_jit(fun, donate="none")
  
          def _fun(x, y, z, *args):
              y = _map_is_array_like(y)
              z = _map_is_array_like(z)
              return fun_jit(x, y, z, *args)
  
          return _fun
  
      def before_filter_jit(fun):
          fun_jit = eqx.filter_jit(fun, donate="none")
  
          def _fun(*args):
              args = _map_is_array_like(args)
              return fun_jit(*args)
  
          return _fun
  
      identity = lambda f: f
      for before_jit in (before_filter_jit, identity):
          for after_jit in (after_filter_jit, identity):
  
              @before_jit
              def f(x):
                  return x + 1
  
              primals, tangents = after_jit(eqx.filter_jvp)(
                  f, (1.0,), (1.0,)
              )
              assert tree_allclose(primals, jnp.array(2.0))
              assert tree_allclose(tangents, jnp.array(1.0))
  
              another_object = object()
  
              @before_jit
              def g(*x):
                  y = (
                      1 + x[0] ** 2,
                      2 + x[1] * 3,
                      3 + x[2] * 4,
                      4 + x[3] * 5,
                      5 + x[4] * 6,
                      another_object,
                  )
                  return y
  
              primals_in = (
                  jnp.array(1.0),
                  jnp.array(1),
                  1.0,
                  1,
                  jnp.array(1.0),
                  object(),
              )
              true_primals_out = (
                  jnp.array(2.0),
                  jnp.array(5),
                  jnp.array(7.0),
                  jnp.array(9),
                  jnp.array(11.0),
                  another_object,
              )
  
              tangents_in1 = (
                  jnp.array(5.0),
                  None,
                  None,
                  None,
                  None,
                  None,
              )
              true_tangents_out1 = (
                  jnp.array(10.0),
                  None,
                  None,
                  None,
                  None,
                  None,
              )
  
              tangents_in2 = (
                  jnp.array(5.0),
                  None,
                  jnp.array(3.0),
                  None,
                  None,
                  None,
              )
              true_tangents_out2 = (
                  jnp.array(10.0),
                  None,
                  jnp.array(12.0),
                  None,
                  None,
                  None,
              )
  
              tangents_in3 = (None, None, None, None, None, None)
              true_tangents_out3 = (None, None, None, None, None, None)
  
              all_tangents_in = (
                  tangents_in1,
                  tangents_in2,
                  tangents_in3,
              )
              all_true_tangents_out = (
                  true_tangents_out1,
                  true_tangents_out2,
                  true_tangents_out3,
              )
  
              for tangents_in, true_tangents_out in zip(
                  all_tangents_in, all_true_tangents_out
              ):
                  primals_out, tangents_out = after_jit(eqx.filter_jvp)(
                      g, primals_in, tangents_in
                  )
                  assert primals_out == true_primals_out
                  assert tree_allclose(tangents_out, true_tangents_out)
  
              bad_tangents_in1 = (
                  jnp.array(5),
                  None,
                  None,
                  None,
                  None,
                  None,
              )
              bad_tangents_in2 = (
                  None,
                  None,
                  None,
                  None,
                  None,
                  object(),
              )
              bad_tangents_in3 = (
                  None,
                  jnp.array(1.0),
                  None,
                  None,
                  None,
                  None,
              )
  
              for tangents_in in (
                  bad_tangents_in1,
                  bad_tangents_in2,
                  bad_tangents_in3,
              ):
                  with pytest.raises(TypeError):
                      after_jit(eqx.filter_jvp)(
                          g, primals_in, tangents_in
                      )
  
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 7 for calling a global function with example usage: test_filter_jvp() and returns: _fun, _fun, fun_jit, fun_jit, unknown, y