code: def extract_pyramid_video_tokens(
      tokens: jnp.ndarray,
      num_pyramid_levels: int,
      feature_pyramid_downsample_stride: int,
      num_video_tokens_level0: int,
      num_text_tokens: int,
  ) -> jnp.ndarray:
      """Extracts video tokens from a feature pyramid, removing text tokens.
  
      Args:
        tokens: Concatenated tokens from each pyramid level in shape (batch,
          seq_len_level0 + seq_len_level1 + ..., channels). This include the text
          tokens.
        num_pyramid_levels: Number of feature pyramid levels.
        feature_pyramid_downsample_stride: Downsample stride used to create the
          feature pyramid.
        num_video_tokens_level0: Number of video tokens in the first level of
          feature pyramid.
        num_text_tokens: Number fo text tokens.
  
      Returns:
        Concatenated video tokens from each pyramid level in shape (batch,
          seq_len_level0 + seq_len_level1 + ..., channels).
      """
  
      num_video_tokens_per_level = [
          (
              num_video_tokens_level0
              // (feature_pyramid_downsample_stride**idx)
          )
          for idx in range(num_pyramid_levels)
      ]
      tokens_per_level = split_pyramid_features(
          tokens,
          num_video_tokens_level0,
          num_pyramid_levels,
          feature_pyramid_downsample_stride,
          num_text_tokens,
          np_backend=jnp,
      )
      video_tokens_per_level = []
      for n, t in zip(num_video_tokens_per_level, tokens_per_level):
          video_tokens_per_level.append(t[:, :n])
      return jnp.concatenate(video_tokens_per_level, 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 1 for calling a global function with example usage: extract_pyramid_video_tokens(tokens, num_pyramid_levels, feature_pyramid_downsample_stride, num_video_tokens_level0, num_text_tokens) and returns: jnp