code: def test_map(self):
      init_fn, _ = transform.transform(get_net)
      params = init_fn(jax.random.PRNGKey(428), jnp.ones((1, 1)))
  
      # parse by layer
      def map_fn(module_name, name, v):
          del name
          if "first_layer" in module_name:
              return v
          else:
              return 2.0 * v
  
      new_params = filtering.map(map_fn, params)
      self.assertLen(jax.tree_util.tree_leaves(new_params), 4)
  
      first_layer_params, second_layer_params = filtering.partition(
          lambda module_name, *_: module_name == "first_layer", params
      )
      for mn in first_layer_params:
          for n in first_layer_params[mn]:
              self.assertEqual(params[mn][n], new_params[mn][n])
  
      for mn in second_layer_params:
          for n in second_layer_params[mn]:
              self.assertEqual(2.0 * params[mn][n], new_params[mn][n])
  
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 11 for calling the FilteringTest obj's underlying member method code with example usage: obj.test_map() and returns: v, unknown