code: def split_pyramid_features(
      tokens: Array,
      num_features_level0: int,
      num_pyramid_levels: int,
      feature_pyramid_downsample_stride: int,
      num_extra_features_per_level: int = 0,
      axis: int = 1,
      np_backend: PyModule = jnp,
  ):
      """Split tokens into a list based on pyramid levels."""
      indices = create_pyramid_split_indices(
          num_features_level0,
          num_pyramid_levels,
          feature_pyramid_downsample_stride,
          num_extra_features_per_level,
      )
      return np_backend.split(tokens, indices, axis=axis)
  
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 3 for calling a global function with example usage: split_pyramid_features(tokens, num_features_level0, num_pyramid_levels, feature_pyramid_downsample_stride, num_extra_features_per_level, axis, np_backend) and returns: np_backend