code: def merge_pyramid_input_masks(
      input_masks: Sequence[jnp.ndarray],
      input_text_mask: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Merges input masks from different pyramid levels into one."""
      if input_text_mask is None:
          return jnp.concatenate(input_masks, axis=1)
      all_masks = []
      for mask in input_masks:
          all_masks.append(
              jnp.concatenate([mask, input_text_mask], axis=1)
          )
      return jnp.concatenate(all_masks, 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 5 for calling a global function with example usage: merge_pyramid_input_masks(input_masks, input_text_mask) and returns: jnp, jnp