code: def topk_mask(value: jax.Array, density_fraction: float) -> jax.Array:
      """Return a mask with 1s marking the top fraction of value.
  
      Note: This routine takes care to make sure that ties are handled without
      bias toward smaller indices.  This can be a problem when pruning large
      embedding matrices, or global pruning where all parameters in the model
      are concatenated together and pruned at once.
  
      Args:
        value: An array. Must contain sortable values (i.e. not complex).
        density_fraction: A float. What fraction of value should be kept.
  
      Returns:
        A mask containing 1s where the topk elements of value are. k is
        determined based on density_fraction and the size of value.
      """
  
      def topk_mask_internal(value):
          assert value.ndim == 1
          indices = jnp.argsort(value)
          k = jnp.round(density_fraction * jnp.size(value)).astype(
              jnp.int32
          )
          mask = jnp.greater_equal(
              np.arange(value.size), value.size - k
          )
          mask = jnp.zeros_like(mask).at[indices].set(mask)
          return mask.astype(np.int32)
  
      # shuffle value so that identical values aren't always pruned
      # with a bias to lower indices
      orig_shape = value.shape
      value = jnp.reshape(value, -1)
      shuffled_indices = jax.random.shuffle(
          jax.random.PRNGKey(42),
          jnp.arange(0, jnp.size(value), dtype=jnp.int32),
      )
  
      shuffled_mask = topk_mask_internal(value[shuffled_indices])
      mask = (
          jnp.zeros_like(shuffled_mask)
          .at[shuffled_indices]
          .set(shuffled_mask)
      )
      mask = jnp.reshape(mask, orig_shape)
      return mask
  
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 1 for calling a global function with example usage: topk_mask(value, density_fraction) and returns: mask, mask