code: def get_tube_mask_indices(
      n_batch: int,
      n_tokens: int,
      token_mask_probability: float,
      temporal_dims: int,
      rng: jax.Array,
  ):
      """Returns indices to use for tube masking in VideoMAE.
  
      The difference between the random and tube masking is that the tube masking
      takes into account the temporal dimension when masking.
  
      Args:
        n_batch: The batch size of the sequence to generate.
        n_tokens: The number of tokens.
        token_mask_probability: Probability of dropping out the input tokens
        during training.
        temporal_dims: The temporal dimension.
        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.
      """
  
      n_tokens_frame = n_tokens // temporal_dims
      n_masked_frame = int(token_mask_probability * n_tokens_frame)
      batch_indices = jnp.arange(n_batch).reshape(n_batch, 1)
  
      mask_indices_frame, _ = shuffle_and_partition(
          n_batch, n_tokens_frame, n_masked_frame, rng
      )
      binary_mask_frame = (
          jnp.zeros((n_batch, n_tokens_frame))
          .at[batch_indices, mask_indices_frame]
          .set(1.0)
      )
  
      # Add temporal dims
      binary_mask = jnp.tile(binary_mask_frame, [1, temporal_dims])
  
      # Apply binary_mask
      n_masked_tokens = n_masked_frame * temporal_dims
      n_unmasked_tokens = n_tokens - n_masked_tokens
      masked_indices = jnp.nonzero(
          binary_mask, size=(n_batch * n_masked_tokens)
      )[1].reshape(n_batch, -1)
      unmasked_indices = jnp.nonzero(
          binary_mask - 1, size=(n_batch * n_unmasked_tokens)
      )[1].reshape(n_batch, -1)
  
      return masked_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 3 for calling a global function with example usage: get_tube_mask_indices(n_batch, n_tokens, token_mask_probability, temporal_dims, rng) and returns: masked_indices,  unmasked_indices,  binary_mask