class_docstr: MTV model.
docstr: Extracts encoder output.
docstr_tok: ["Extracts","encoder","output","."]
code: def _extract_encoder_output(
      self, x: jnp.ndarray, axis: int = 1
  ) -> jnp.ndarray:
      """Extracts encoder output."""
      if self.classifier in ["token", "0"]:
          x = x.take(indices=0, axis=axis)
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          x = fn(x, axis=list(range(axis, x.ndim - 1)))
      return x
  
imports: import functools
  from typing import Any, List, Sequence, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.mtv import model_utils
  from scenic.projects.vivit import model
  from scenic.train_lib import train_utils
label_desc: Use image node_id 4 for calling the MTV obj's underlying member method code with example usage: obj._extract_encoder_output(x, axis) and returns: x