code: def test_filter_vjp(getkey):
      mlp = eqx.nn.MLP(2, 3, 2, 2, key=getkey())
      x = jnp.array(1)
      y = 1.0
      sentinel = object()
      sentinel2 = object()
  
      def f(_mlp, _x, _y, _sentinel):
          return _mlp(jnp.array([1.0, 2.0])) + _x + _y, 2.0, sentinel2
  
      def g(_mlp, _x, _y, _sentinel):
          return f(_mlp, _x, _y, _sentinel), _sentinel
  
      def _check(x, y):
          assert x.shape == y.shape
          assert x.dtype == y.dtype
  
      out1, vjpfun1 = eqx.filter_vjp(f, mlp, x, y, sentinel)
      out2, vjpfun2, aux = eqx.filter_vjp(
          g, mlp, x, y, sentinel, has_aux=True
      )
      assert aux is sentinel
      for out, vjpfun in ((out1, vjpfun1), (out2, vjpfun2)):
          out_array, out_float, out_sentinel = out
          assert out_array.shape == (3,)
          assert out_array.dtype == jnp.float32
          assert tree_allclose(out_float, 2.0)
          assert out_sentinel is sentinel2
          ct_mlp, ct_x, ct_y, ct_sentinel = vjpfun(
              (jnp.array([1.0, 2.0, 3.0]), None, None)
          )
          mlp_dyn = eqx.filter(mlp, eqx.is_array)
          assert jtu.tree_structure(mlp_dyn) == jtu.tree_structure(
              ct_mlp
          )
          jtu.tree_map(_check, mlp_dyn, ct_mlp)
          assert ct_y is None
          assert ct_sentinel is None
  
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 8 for calling a global function with example usage: test_filter_vjp(getkey) and returns: unknown,  int,  sentinel2, f,  _sentinel