code: def test_filter(getkey):
      filter_fn = lambda x: isinstance(x, int)
      for pytree in (
          [
              1,
              2,
              [
                  3,
                  "hi",
                  {
                      "a": jnp.array(1),
                      "b": 4,
                      "c": eqx.nn.MLP(2, 2, 2, 2, key=getkey()),
                  },
              ],
          ],
          [1, 1, 1, 1, "hi"],
      ):
          filtered = eqx.filter(pytree, filter_spec=filter_fn)
          for arg in jtu.tree_leaves(filtered):
              assert isinstance(arg, int)
          num_int_leaves = sum(
              1
              for leaf in jtu.tree_leaves(filtered)
              if isinstance(leaf, int)
          )
          assert len(jtu.tree_leaves(filtered)) == num_int_leaves
  
      filter_spec = [False, True, [filter_fn, True]]
      sentinel = object()
      pytree = [
          eqx.nn.Linear(1, 1, key=getkey()),
          eqx.nn.Linear(1, 1, key=getkey()),
          [eqx.nn.Linear(1, 1, key=getkey()), sentinel],
      ]
      filtered = eqx.filter(pytree, filter_spec=filter_spec)
      none_linear = jtu.tree_map(
          lambda _: None, eqx.nn.Linear(1, 1, key=getkey())
      )
      assert filtered[0] == none_linear
      assert filtered[1] == pytree[1]
      assert filtered[2][0] == none_linear
      assert filtered[2][1] is sentinel
  
      with pytest.raises(ValueError):
          eqx.filter(pytree, filter_spec=filter_spec[1:])
  
file_path: equinox/tests/test_filters.py
imports: from typing import Any
  import equinox
  import jax.numpy
  import jax.tree_util
  import numpy
  import pytest
label_desc: Use image node_id 5 for calling a global function with example usage: test_filter(getkey) without return types