code: def test_apply_update(self):
      # test list and tuple
      def update_fn(path, value):
          """Multiply the first element of all leaf nodes of the pytree by -1."""
          if path[-1] == "0" and isinstance(value, int):
              return value * -1
          return value
  
      for tuple_wrap in (lambda x: x, tuple):
          l = tuple_wrap([tuple_wrap([1, 2]), tuple_wrap([3, 4])])
          c = cursor(l)
          l2 = c.apply_update(update_fn).build()
          self.assertEqual(
              l2, tuple_wrap([tuple_wrap([-1, 2]), tuple_wrap([-3, 4])])
          )
          self.assertEqual(
              l, tuple_wrap([tuple_wrap([1, 2]), tuple_wrap([3, 4])])
          )  # make sure the original object is unchanged
  
      # test regular dict and FrozenDict
      def update_fn(path, value):
          """Multiply all dense kernel params by 2 and add 1.
          Subtract the Dense_1 bias param by 1."""
          if "kernel" in path:
              return value * 2 + 1
          elif "Dense_1" in path and "bias" in path:
              return value - 1
          return value
  
      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.key(0), jnp.empty((1, 2)))[
                  "params"
              ]
          )
  
          c = cursor(params)
          params2 = c.apply_update(update_fn).build()
          for layer in ("Dense_0", "Dense_1", "Dense_2"):
              self.assertTrue(
                  (
                      params2[layer]["kernel"]
                      == 2 * params[layer]["kernel"] + 1
                  ).all()
              )
              if layer == "Dense_1":
                  self.assertTrue(
                      (
                          params2[layer]["bias"]
                          == jnp.array([-1, -1, -1])
                      ).all()
                  )
              else:
                  self.assertTrue(
                      (
                          params2[layer]["bias"]
                          == params[layer]["bias"]
                      ).all()
                  )
          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.key(0), jnp.empty((1, 2))
                          )["params"]
                      ),
                  )
              )
          )  # make sure original params are unchanged
  
      # test TrainState
      def update_fn(path, value):
          """Replace params with empty dictionary."""
          if "params" in path:
              return {}
          return value
  
      state = train_state.TrainState.create(
          apply_fn=lambda x: x,
          params={"a": 1, "b": 2},
          tx=optax.adam(1e-3),
      )
      c = cursor(state)
      state2 = c.apply_update(update_fn).build()
      self.assertEqual(state2.params, {})
      self.assertEqual(
          state.params, {"a": 1, "b": 2}
      )  # make sure original params are unchanged
  
      # test NamedTuple
      def update_fn(path, value):
          """Add 5 to all x-attribute values that are ints."""
          if path[-1] == "x" and isinstance(value, int):
              return value + 5
          return value
  
      t = GenericTuple(GenericTuple(0, 1), GenericTuple(2, 3))
      c = cursor(t)
      t2 = c.apply_update(update_fn).build()
      self.assertEqual(
          t2, GenericTuple(GenericTuple(5, 1), GenericTuple(7, 3))
      )
      self.assertEqual(
          t, GenericTuple(GenericTuple(0, 1), GenericTuple(2, 3))
      )
  
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 6 for calling the CursorTest obj's underlying member method code with example usage: obj.test_apply_update() and returns: value, value, value, value, unknown, unknown, x, dict, unknown, unknown