code: def test_find_all(self):
      # test list and tuple
      def cond_fn(path, value):
          """Get all lists that are not the first element in its parent."""
          return path[-1] != "0" and isinstance(value, (tuple, list))
  
      for tuple_wrap in (lambda x: x, tuple):
          l = tuple_wrap(
              [
                  tuple_wrap([1, 2]),
                  tuple_wrap([3, 4]),
                  tuple_wrap([5, 6]),
              ]
          )
          c = cursor(l)
          c2, c3 = c.find_all(cond_fn)
          c2[0] *= -1
          c3[1] *= -2
          self.assertEqual(
              c.build(),
              tuple_wrap(
                  [
                      tuple_wrap([1, 2]),
                      tuple_wrap([-3, 4]),
                      tuple_wrap([5, -12]),
                  ]
              ),
          )
          self.assertEqual(
              l,
              tuple_wrap(
                  [
                      tuple_wrap([1, 2]),
                      tuple_wrap([3, 4]),
                      tuple_wrap([5, 6]),
                  ]
              ),
          )  # make sure the original object is unchanged
  
      # test regular dict and FrozenDict
      def cond_fn(path, value):
          """Get the second and third dense params."""
          return "Dense_1" in path or "Dense_2" in path
  
      class Model(nn.Module):
          @nn.compact
          def __call__(self, x):
              x = nn.Dense(3)(x)
              x = nn.relu(x)
              x = nn.Dense(3)(x)
              x = nn.relu(x)
              x = nn.Dense(3)(x)
              x = nn.relu(x)
              return x
  
      for freeze_wrap in (lambda x: x, freeze):
          params = freeze_wrap(
              Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))[
                  "params"
              ]
          )
          c = cursor(params)
          for i, c2 in enumerate(c.find_all(cond_fn)):
              self.assertEqual(
                  c2["kernel"].set(123)[f"Dense_{i+1}"],
                  freeze_wrap(
                      {
                          "kernel": 123,
                          "bias": params[f"Dense_{i+1}"]["bias"],
                      }
                  ),
              )
          self.assertTrue(
              jax.tree_util.tree_all(
                  jax.tree_util.tree_map(
                      lambda x, y: (x == y).all(),
                      params,
                      freeze_wrap(
                          Model().init(
                              jax.random.PRNGKey(0), jnp.empty((1, 2))
                          )["params"]
                      ),
                  )
              )
          )  # make sure original params are unchanged
  
      # test TrainState
      def cond_fn(path, value):
          """Find TrainState params."""
          return "params" in path
  
      state = train_state.TrainState.create(
          apply_fn=lambda x: x,
          params={"a": 1, "b": 2},
          tx=optax.adam(1e-3),
      )
      c = cursor(state)
      c2 = list(c.find_all(cond_fn))
      self.assertEqual(len(c2), 1)
      c2 = c2[0]
      self.assertEqual(c2["b"].set(-1).params, {"a": 1, "b": -1})
      self.assertEqual(
          state.params, {"a": 1, "b": 2}
      )  # make sure original params are unchanged
  
      # test NamedTuple
      def cond_fn(path, value):
          """Get all GenericTuples that have int x-attribute values."""
          return isinstance(value, GenericTuple) and isinstance(
              value.x, int
          )
  
      t = GenericTuple(
          GenericTuple(0, "a"),
          GenericTuple(1, "b"),
          GenericTuple("c", 2),
      )
      c = cursor(t)
      c2, c3 = c.find_all(cond_fn)
      c2.x += 5
      c3.x += 6
      self.assertEqual(
          c.build(),
          GenericTuple(
              GenericTuple(5, "a"),
              GenericTuple(7, "b"),
              GenericTuple("c", 2),
          ),
      )
      self.assertEqual(
          t,
          GenericTuple(
              GenericTuple(0, "a"),
              GenericTuple(1, "b"),
              GenericTuple("c", 2),
          ),
      )
  
imports: import dataclasses
  from typing import Any, NamedTuple
  import jax
  import jax.numpy
  import optax
  from absl.testing import absltest
  import flax
  import flax.linen
  from flax.core import freeze
  from flax.cursor import AccessType, _traverse_tree, cursor
  from flax.errors import CursorFindError, TraverseTreeError
  from flax.training import train_state
label_desc: Use image node_id 12 for calling the CursorTest obj's underlying member method code with example usage: obj.test_find_all() and returns: unknown, unknown, unknown, unknown, x