code: def create_pyramid_input_masks(
      input_mask: jnp.ndarray,
      num_features_level0: int,
      num_pyramid_levels: int,
      feature_pyramid_downsample_stride: int,
      num_text_tokens: int,
  ) -> List[jnp.ndarray]:
      """Splits input mask into a list of them based on pyramid levels."""
      split_indices = create_pyramid_split_indices(
          num_features_level0,
          num_pyramid_levels,
          feature_pyramid_downsample_stride,
          num_extra_features_per_level=num_text_tokens,
      )
      return np.split(input_mask, split_indices, axis=1)
  
imports: from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
  from absl import logging
  import flax
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.vivit import model_utils
  from scenic.train_lib import train_utils
label_desc: Use image node_id 4 for calling a global function with example usage: create_pyramid_input_masks(input_mask, num_features_level0, num_pyramid_levels, feature_pyramid_downsample_stride, num_text_tokens) and returns: np