code: def test_filter_grad(getkey):
      a = jrandom.normal(getkey(), (2, 3))
      b = jrandom.normal(getkey(), (2, 3))
  
      @eqx.filter_grad
      def f(x):
          sum = jnp.array(0.0)
          for arg in jtu.tree_leaves(x):
              if eqx.is_array_like(arg):
                  sum = sum + jnp.sum(arg).astype(sum.dtype)
          return sum
  
      ga, gb = f([a, b])
      assert tree_allclose(ga, jnp.ones((2, 3)))
      assert tree_allclose(gb, jnp.ones((2, 3)))
  
      gtrue, ghi, gobject, ga = f([True, "hi", object(), a])
      assert gtrue is None
      assert ghi is None
      assert gobject is None
      assert tree_allclose(ga, jnp.ones((2, 3)))
  
      gtrue, gdict, (g5, g1), gnp = f(
          [
              True,
              {"hi": eqx.nn.Linear(1, 1, key=getkey())},
              (5, 1.0),
              np.array([2.0, 3.0]),
          ]
      )
      assert gtrue is None
      assert list(gdict.keys()) == ["hi"]
      assert isinstance(gdict["hi"], eqx.nn.Linear)
      assert tree_allclose(gdict["hi"].weight, jnp.ones((1, 1)))
      assert tree_allclose(gdict["hi"].bias, jnp.ones(1))
      assert g5 is None
      assert g1 is None
      assert tree_allclose(gnp, jnp.ones(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 1 for calling a global function with example usage: test_filter_grad(getkey) and returns: sum