code: def test_partition_and_combine(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, unfiltered = eqx.partition(
              pytree, filter_spec=filter_fn
          )
          for arg in jtu.tree_leaves(filtered):
              assert isinstance(arg, int)
          for arg in jtu.tree_leaves(unfiltered):
              assert not isinstance(arg, int)
          assert eqx.combine(filtered, unfiltered) == pytree
          assert eqx.combine(unfiltered, filtered) == pytree
  
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 6 for calling a global function with example usage: test_partition_and_combine(getkey) without return types