code: def test_filter_custom_jvp_kwargs():
      @eqx.filter_custom_jvp
      def call(x, y, *, fn):
          return fn(x, y)
  
      was_called = False
  
      @call.def_jvp
      def call_jvp(primals, tangents, *, fn):
          nonlocal was_called
          was_called = True
          x, y = primals
          tx, ty = tangents
          primal_out = call(x, y, fn=fn)
          tangent_out = tx**2 + ty
          return primal_out, tangent_out
  
      f = lambda a, b: a**2 + b
      assert tree_allclose(call(1, 2, fn=f), 3)
      assert tree_allclose(call(1.0, 2.0, fn=f), 3.0)
      assert tree_allclose(call(jnp.array(1), 2, fn=f), jnp.array(3))
      assert tree_allclose(
          call(jnp.array(1.0), 2, fn=f), jnp.array(3.0)
      )
  
      def jvpcall(x, y, tx, ty):
          def _jvpcall(_x, _y):
              return call(_x, _y, fn=f)
  
          return jax.jvp(_jvpcall, (x, y), (tx, ty))
  
      primal_out, tangent_out = jvpcall(2.0, 1.5, 3.0, 4.0)
      assert was_called
      assert tree_allclose(primal_out, 5.5)
      assert tree_allclose(tangent_out, 13.0)
  
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 13 for calling a global function with example usage: test_filter_custom_jvp_kwargs() and returns: fn, primal_out,  tangent_out, jax, call