code: def test_param_selection(self):
      params = {
          "x": {
              "kernel": 1,
              "bias": 2,
              "y": {
                  "kernel": 3,
                  "bias": 4,
              },
              "z": {},
          },
      }
      expected_params = {
          "x": {
              "kernel": 2,
              "bias": 2,
              "y": {
                  "kernel": 6,
                  "bias": 4,
              },
              "z": {},
          },
      }
      names = []
  
      def filter_fn(name, _):
          names.append(
              name
          )  # track names passed to filter_fn for testing
          return "kernel" in name
  
      traversal = traverse_util.ModelParamTraversal(filter_fn)
  
      values = list(traversal.iterate(params))
      configs = [
          (params, expected_params),
          (
              flax.core.FrozenDict(params),
              flax.core.FrozenDict(expected_params),
          ),
      ]
      for model, expected_model in configs:
          self.assertEqual(values, [1, 3])
          self.assertEqual(
              set(names),
              set(["/x/kernel", "/x/bias", "/x/y/kernel", "/x/y/bias"]),
          )
          new_model = traversal.update(lambda x: x + x, model)
          self.assertEqual(new_model, expected_model)
  
imports: import collections
  import jax
  import jax.numpy
  import numpy
  import optax
  from absl.testing import absltest
  import flax
  from flax import traverse_util
  from flax.core import freeze
label_desc: Use image node_id 2 for calling the ModelParamTraversalTest obj's underlying member method code with example usage: obj.test_param_selection() and returns: unknown