code: def get_mask_indices(
      n_batch: int, n_tokens: int, n_masked: int, rng: jax.Array
  ):
      """Returns indices to use for masking in MAE.
  
      Args:
        n_batch: The batch size of the sequence to generate.
        n_tokens: The number of tokens.
        n_masked: The number of tokens to mask. Must have 0 <= n_masked < n_tokens.
        rng: The random number key.
  
      Returns:
        Three arrays. masked_indices of shape [n_batch, n_masked], unmasked_indices
        of shape [n_batch, n_tokens - n_masked] and binary_mask of shape
        [n_batch, n_tokens] where 1 indicates that the token is masked.
      """
      batch_indices = jnp.arange(n_batch).reshape(n_batch, 1)
      mask_indices, unmasked_indices = shuffle_and_partition(
          n_batch, n_tokens, n_masked, rng
      )
      binary_mask = (
          jnp.zeros((n_batch, n_tokens))
          .at[batch_indices, mask_indices]
          .set(1.0)
      )
  
      return mask_indices, unmasked_indices, binary_mask
  
file_path: scenic/scenic/projects/av_mae/model_utils.py
imports: from typing import Callable, Iterable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 2 for calling a global function with example usage: get_mask_indices(n_batch, n_tokens, n_masked, rng) and returns: mask_indices,  unmasked_indices,  binary_mask