code: def test_partition(self):
      init_fn, _ = transform.transform(get_net)
      params = init_fn(jax.random.PRNGKey(428), jnp.ones((1, 1)))
  
      # parse by layer
      first_layer_params, second_layer_params = filtering.partition(
          lambda module_name, *_: module_name == "first_layer", params
      )
      self.assertEqual(
          get_names(first_layer_params),
          {"first_layer/w", "first_layer/b"},
      )
      self.assertEqual(
          get_names(second_layer_params),
          {"second_layer/w", "second_layer/b"},
      )
  
      # parse by variable type
      weights, biases = filtering.partition(
          lambda module_name, name, _: name == "w", params
      )  # pytype: disable=wrong-arg-types
      self.assertEqual(
          get_names(weights), {"first_layer/w", "second_layer/w"}
      )
      self.assertEqual(
          get_names(biases), {"first_layer/b", "second_layer/b"}
      )
  
      # Compose regexes
      regex = compile_regex(["first_layer.*", ".*w"])
      matching, not_matching = filtering.partition(
          lambda module_name, name, _: regex.match(
              f"{module_name}/{name}"
          ),
          params,
      )
      self.assertEqual(
          get_names(matching),
          {"first_layer/w", "first_layer/b", "second_layer/w"},
      )
      self.assertEqual(get_names(not_matching), {"second_layer/b"})
  
      matching, not_matching = filtering.partition(
          lambda mod_name, name, _: mod_name == "first_layer"
          and name != "w",
          params,
      )
      self.assertEqual(get_names(matching), {"first_layer/b"})
      self.assertEqual(
          get_names(not_matching),
          {"first_layer/w", "second_layer/w", "second_layer/b"},
      )
  
imports: import collections
  from collections.abc import Sequence
  import itertools
  import re
  import types
  from typing import Any, Callable
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import basic
  from haiku._src import data_structures
  from haiku._src import filtering
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the FilteringTest obj's underlying member method code with example usage: obj.test_partition() without return types