code: def prepend_cls(x: jnp.ndarray, cls: jnp.ndarray) -> jnp.ndarray:
      """Prepend a CLS token."""
      assert x.ndim == 3 and cls.ndim == 3
      n, _, _ = x.shape
      cls = jnp.tile(cls, [n, 1, 1])
      return jnp.concatenate([cls, x], axis=1)
  
imports: import functools
  from typing import Callable, List, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
  from scenic.projects.unloc import encoders
  from scenic.projects.unloc import model_utils
label_desc: Use image node_id 1 for calling a global function with example usage: prepend_cls(x, cls) and returns: jnp