code: def test_path_aware_map_with_masked(self):
      params = {
          "linear_1": {"w": jnp.zeros((5, 6)), "b": jnp.zeros(5)},
          "linear_2": {"w": jnp.zeros((6, 1)), "b": jnp.zeros(1)},
      }
      gradients = jax.tree_util.tree_map(
          jnp.ones_like, params
      )  # dummy gradients
  
      params_mask = traverse_util.path_aware_map(
          lambda path, x: "w" in path, params
      )
      tx = optax.masked(optax.sgd(1.0), params_mask)
      state = tx.init(params)
      updates, new_state = tx.update(gradients, state, params)
      new_params = optax.apply_updates(params, updates)
  
      self.assertTrue(
          np.allclose(
              new_params["linear_1"]["b"], gradients["linear_1"]["b"]
          )
      )
      self.assertTrue(
          np.allclose(
              new_params["linear_2"]["b"], gradients["linear_2"]["b"]
          )
      )
      self.assertTrue(
          np.allclose(
              new_params["linear_1"]["w"], -gradients["linear_1"]["w"]
          )
      )
      self.assertTrue(
          np.allclose(
              new_params["linear_2"]["w"], -gradients["linear_2"]["w"]
          )
      )
  
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 5 for calling the ModelParamTraversalTest obj's underlying member method code with example usage: obj.test_path_aware_map_with_masked() without return types