code: def get_bottleneck_representation(x, pooling_type):
      """Bottleneck representation.
  
      Args:
          x: input tensor.
          pooling_type: type of the classifier layer. Options are 'gap', 'gmp',
          'gsp', 'token'.
  
      Returns:
          bottleneck tensor.
      """
  
      if pooling_type in ("token", "0"):
          x = x[:, 0]
      elif pooling_type in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              pooling_type
          ]
          x = fn(x, axis=1)
  
          x = jnp.mean(x, axis=list(range(1, x.ndim - 1)))
      else:
          raise ValueError(
              "Pooling type should be in ['gap', 'gmp', 'gsp', 'token']."
          )
  
      return nn_layers.IdentityLayer(name="bottleneck")(x)
  
imports: import functools
  from typing import Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  import numpy
  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.polyvit import polyvit_base_model
  from scenic.projects.vivit import model
label_desc: Use image node_id 2 for calling a global function with example usage: get_bottleneck_representation(x, pooling_type) and returns: nn_layers