code: def token_indexes_not_to_drop(
      seqlen, n_tokens, seqlen_selection, rng
  ):
      """Returns only the token indexes to keep in a sequence of tokens."""
      idx_kept_tokens = jnp.arange(n_tokens)
      if seqlen > 0 and seqlen <= n_tokens:
          if seqlen_selection in ["consecutive", "first"]:
              if seqlen_selection == "first":
                  offset = 0
              else:
                  offset = jax.random.randint(
                      rng, (1,), 0, n_tokens - seqlen + 1
                  )[0]
              # Workaround because jnp.arange(offset, offset + seqlen) causes
              # a ConcretizationError (even though shape is known to be seqlen...)
              idx_kept_tokens = jnp.ones(seqlen) * offset + jnp.arange(
                  seqlen
              )
          elif seqlen_selection == "unstructured":
              idx_kept_tokens = jax.random.permutation(rng, n_tokens)[
                  :seqlen
              ]
      idx_kept_tokens = jnp.asarray(idx_kept_tokens, dtype=jnp.int32)
      return idx_kept_tokens
  
imports: import copy
  import functools
  from typing import Any, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling a global function with example usage: token_indexes_not_to_drop(seqlen, n_tokens, seqlen_selection, rng) and returns: idx_kept_tokens