code: def extract_emb(
      x: jnp.ndarray,
      classifier: str,
      keepdims: bool = False,
      input_mask: Optional[jnp.ndarray] = None,
      input_word_ids: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Extracts an embedding from a sequence of tokens.
  
      Args:
        x: A 3D float tensor of shape (batch_size, sequence_length, channels)
          representing the tokens where we extract the embedding. If x is 2D, we
          return as it is.
        classifier: 'token', 'eos', or 'gap'.
        keepdims: Whether or not to make the output embedding as the same dimension
          as the input tokens.
        input_mask: Optional. A 2D binary tensor of shape (batch_size,
          sequence_length) representing the input mask. This arg is only used when
          classifier = 'gap'.
        input_word_ids: Optional. A 2D binary tensor of shape (batch_size,
          sequence_length) representing the input word ids. This arg is only used
          when classifier = 'eos'. 'eos' tokens are assumed to have the largest ids.
  
      Returns:
        A 3D float tensor of shape (batch_size, 1, channels) if
        keepdims=True or a 2D float tensor of shape (batch_size, channels) if
        keepdims=False.
      """
      if x.ndim == 2:
          logging.info("Input is 2D and return as it is.")
          return x
      if classifier == "token":
          return x[:, :1] if keepdims else x[:, 0]
      elif classifier == "eos" and input_word_ids is not None:
          x = x[jnp.arange(x.shape[0]), input_word_ids.argmax(-1)]
          if keepdims:
              x = jnp.expand_dims(x, axis=1)
          return x
      else:
          if input_mask is not None:
              return jnp.sum(
                  jnp.multiply(x, jnp.expand_dims(input_mask, axis=-1)),
                  axis=1,
                  keepdims=keepdims,
              ) / jnp.expand_dims(
                  input_mask.sum(axis=-1, keepdims=keepdims), axis=-1
              )
          return jnp.mean(x, axis=1, keepdims=keepdims)
  
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 6 for calling a global function with example usage: extract_emb(x, classifier, keepdims, input_mask, input_word_ids) and returns: x, unknown, x, jnp, unknown