code: def shuffle_and_partition(
      n_batch: int, n_tokens: int, n_masked: int, rng: jax.Array
  ):
      """Implements random shuffling and partitioning necessary for 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:
        Two arrays. The first one contains indices of masked tokens, and has
        shape [n_batch, n_masked]. The second contains indices of unmasked tokens
        and has shape [n_batch, n_tokens - n_masked].
      """
      if n_masked >= n_tokens or n_masked < 0:
          raise ValueError(
              f"n_masked = {n_masked} should be >=0 and <{n_tokens}."
          )
  
      ids = jnp.tile(jnp.arange(n_tokens), n_batch).reshape(
          (n_batch, n_tokens)
      )
      n_remainder = n_tokens - n_masked
      if n_masked > 0:
          rng_keys = jax.random.split(rng, n_batch)
          ids = jax.vmap(
              lambda seq, rng: jax.random.permutation(
                  rng, seq, independent=True
              )
          )(ids, rng_keys)
      masked = jax.lax.dynamic_slice(
          ids,
          (
              0,
              0,
          ),
          (
              n_batch,
              n_masked,
          ),
      )
      unmasked = jax.lax.dynamic_slice(
          ids,
          (
              0,
              n_masked,
          ),
          (
              n_batch,
              n_remainder,
          ),
      )
      return masked, unmasked
  
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 1 for calling a global function with example usage: shuffle_and_partition(n_batch, n_tokens, n_masked, rng) and returns: masked,  unmasked