code: def prepend_one(input_mask: jnp.ndarray) -> jnp.ndarray:
      """Prepends ones to the input mask.
  
      Args:
        input_mask: Mask assumed to be of shape [batch, tokens, ...].
  
      Returns:
        Input mask prepended with one in shape [batch, 1+tokens, ...].
      """
      return jnp.concatenate(
          [
              jnp.ones(
                  (input_mask.shape[0], 1), dtype=input_mask.dtype
              ),
              input_mask,
          ],
          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 3 for calling a global function with example usage: prepend_one(input_mask) and returns: jnp