code: def apply_mask(
      params: hk.Params,
      masks: Sequence[hk.Params],
      module_sparsity: ModuleSparsity,
  ) -> hk.Params:
      """Apply existing masks to params based on sparsity_predicate_map.
  
      Some of params may not be masked depending on the content of
      module_sparsity.  masks must have the same structure as implied by
      module_sparsity.
  
      Args:
        params: Tree to mask, can be a superset of masks.
        masks: Tree of masks to apply to params.  This must match the result of
          applying module_sparsity to params.
        module_sparsity: A dictionary mapping predicates to sparsity levels. Any
          leaf matching a predicate key will be pruned to the resulting sparsity
          level.
  
      Returns:
        A tree of masked params.
      """
      params_to_prune, _, params_no_prune = _create_partitions(
          module_sparsity, params
      )
      pruned_params = []
      for value, mask in zip(params_to_prune, masks):
          pruned_params.append(
              jax.tree_util.tree_map(lambda x, y: x * y, value, mask)
          )
      params = hk.data_structures.merge(*pruned_params, params_no_prune)
      return params
  
imports: from collections.abc import Iterator, Mapping, Sequence
  import functools
  from typing import Callable
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 5 for calling a global function with example usage: apply_mask(params, masks, module_sparsity) and returns: params