code: def leastk_mask(scores, ones_fraction):
      """Given a tensor of scores creates a binary mask.
  
      Args:
        scores: top-scores are kept
        ones_fraction: float, of the generated mask.
  
      Returns:
        array, same shape and type as scores or None.
      """
      if ones_fraction is None or ones_fraction == 0:
          return jnp.zeros_like(scores)
      # This is to ensure indices with smallest values are selected.
      scores = -scores
  
      n_ones = jnp.round(jnp.size(scores) * ones_fraction)
      k = jnp.maximum(1, n_ones).astype(int)
      flat_scores = jnp.reshape(scores, -1)
      threshold = jax.lax.sort(flat_scores)[-k]
  
      mask = (flat_scores >= threshold).astype(flat_scores.dtype)
      return mask.reshape(scores.shape)
  
imports: import functools
  import logging
  import flax
  from flax import linen
  import gin
  import jax
  from jax import random
  import jax.numpy
  import optax
label_desc: Use image node_id 1 for calling a global function with example usage: leastk_mask(scores, ones_fraction) and returns: mask, jnp