class_docstr: Projection Head and Classifier for MBT.
  
  Attributes:
    inputs: nd-array, Input data
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token', 'onetoken'
    num_classes: Number of output classes.
    modality_fusion: Tuple with modalities to combine.
    representation_size: Size of the representation layer in the model's head.
      if None, we skip the extra projection + tanh activation at the end.
    num_layers: Number of extra projection + tanh activation layers before
    projection. Must set representation_size.
    return_prelogits: If true, return the final representation of the network
      before the classification head. Useful when using features for a
      downstream task.
    dtype: JAX data type for activations.
class_docstr_tok: ["Projection","Head","and","Classifier","for","MBT",".","Attributes",":","inputs",":","nd-array",",","Input","data","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",",","'onetoken'","num_classes",":","Number","of","output","classes",".","modality_fusion",":","Tuple","with","modalities","to","combine",".","representation_size",":","Size","of","the","representation","layer","in","the","model","'s","head",".","if","None",",","we","skip","the","extra","projection","+","tanh","activation","at","the","end",".","num_layers",":","Number","of","extra","projection","+","tanh","activation","layers","before","projection",".","Must","set","representation_size",".","return_prelogits",":","If","true",",","return","the","final","representation","of","the","network","before","the","classification","head",".","Useful","when","using","features","for","a","downstream","task",".","dtype",":","JAX","data","type","for","activations","."]
docstr: Applies projection and classifier on the inputs.
code: def __call__(
      self,
      x: Dict[str, Any],
      temporal_dims: Dict[str, Any],
      *,
      train: bool,
  ):
      """Applies projection and classifier on the inputs."""
      if self.num_layers > 1:
          assert (
              self.representation_size is not None
          ), "Please provide representation_size"
      x_out = {}
      counter = 0
      if self.classifier in ["onetoken", "token", "0"]:
          # Obtaining the CLS tokens for each modality.
          # Note when self.classifier is 'onetoken', counter remains 0.
          for modality in self.modality_fusion:
              x_out[modality] = x[:, counter]
              counter += temporal_dims[modality] + 1
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          # Note here we pool each modality separately
          for modality in self.modality_fusion:
              modality_tokens = x[
                  :, counter : counter + temporal_dims[modality]
              ]
              x_out[modality] = fn(
                  modality_tokens,
                  axis=list(range(1, modality_tokens.ndim - 1)),
              )
              counter += temporal_dims[modality]
  
      if self.representation_size is not None:
          for layer in range(self.num_layers):
              if (
                  self.num_layers == 1
              ):  # backward compatibility with previous models
                  name = "pre_logits"
              else:
                  name = "pre_logits_fc_{}".format(layer)
              pre_logits_fc = nn.Dense(
                  self.representation_size, name=name
              )
              if isinstance(x_out, dict):
                  for modality in x_out:
                      x_out[modality] = pre_logits_fc(x_out[modality])
                      x_out[modality] = nn.tanh(x_out[modality])
              else:
                  x_out = pre_logits_fc(x_out)
                  x_out = nn.tanh(x_out)
      else:
          if not isinstance(x_out, dict):
              x_out = nn_layers.IdentityLayer(name="pre_logits")(x_out)
  
      if self.return_prelogits:
          return x_out
      if isinstance(x_out, dict):
          output_projection_fc = nn.Dense(
              self.num_classes,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )
          x_pool = 0
          for modality in x_out:
              x_out[modality] = output_projection_fc(x_out[modality])
              x_pool += x_out[modality]
          x_pool /= len(x_out)
          # We always use the average CLS logits during inference.
          if not train:
              return x_pool
      else:
          x_out = nn.Dense(
              self.num_classes,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )(x_out)
          logging.info("Shape of final logits is %s", x_out.shape)
      return x_out
  
imports: import functools
  import re
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import model_utils
  from scenic.projects.baselines import vit
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model_utils
  import scipy
label_desc: Use image node_id 1 for calling the ProjectionHeadAndClassifier obj's underlying member method code with example usage: obj.__call__(x, temporal_dims) and returns: x_out, x_out, x_pool