code: def update_mask(
      params: hk.Params,
      sparsity_fraction: float,
      module_sparsity: ModuleSparsity,
  ) -> Sequence[hk.Params]:
      """Generate masks based on module_sparsity and sparsity_fraction."""
      params_to_prune, sparsities, _ = _create_partitions(
          module_sparsity, params
      )
      masks = []
  
      def map_fn(x: jax.Array, sparsity: float) -> jax.Array:
          return topk_mask(
              jnp.abs(x), 1.0 - sparsity * sparsity_fraction
          )
  
      for tree, sparsity in zip(params_to_prune, sparsities):
          map_fn_sparsity = functools.partial(map_fn, sparsity=sparsity)
          mask = jax.tree_util.tree_map(map_fn_sparsity, tree)
          masks.append(mask)
      return masks
  
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 6 for calling a global function with example usage: update_mask(params, sparsity_fraction, module_sparsity) and returns: masks, topk_mask