class_docstr: Audio-Visual Fusion Transformer model for Video.
  
  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_classes: Number of output classes.
    modality_fusion: Tuple with modalities to combine.
    fusion_layer: Which layer to fuse modalities.
    num_heads: Number of self-attention heads.
    num_layers: Number of layers.
    patches: Configuration of the patches extracted in the stem of the model.
    hidden_size: Size of the hidden state of the output of model's stem.
    representation_size: Size of the representation layer in the model's head.
      if None, we skip the extra projection + tanh activation at the end.
    temporal_encoding_config: ConfigDict which defines the type of input
      encoding when tokenising the video.
    attention_config: ConfigDict which defines the type of spatio-temporal
      attention applied in the model.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    stochastic_droplayer_rate: Probability of dropping a layer. Linearly
      increases from 0 to the provided value..
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token'.
    return_prelogits: If true, return the final representation of the network
      before the classification head. Useful when using features for a
      downstream task.
    return_preclassifier: If true, return a dict of all token embeddings.
      Useful when using token embeddings for a downstream task.
    return_as_dict: If true, return the token embeddings as a dictionary instead
      of a concatenated tensor.
    use_bottleneck: If True, adds self-attention bottleneck.
    n_bottlenecks: Number of bottleneck tokens.
    test_with_bottlenecks: Whether to use bottlenecks at test time.
    share_encoder: If True, different modalities share the same encoder weights
      for the layers before fusion.
    return_bottlenecks: If True, return bottleneck embeddings.
    use_modality_tokens: If True, modality tokens are used.
    dtype: JAX data type for activations.
class_docstr_tok: ["Audio-Visual","Fusion","Transformer","model","for","Video",".","Attributes",":","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_classes",":","Number","of","output","classes",".","modality_fusion",":","Tuple","with","modalities","to","combine",".","fusion_layer",":","Which","layer","to","fuse","modalities",".","num_heads",":","Number","of","self-attention","heads",".","num_layers",":","Number","of","layers",".","patches",":","Configuration","of","the","patches","extracted","in","the","stem","of","the","model",".","hidden_size",":","Size","of","the","hidden","state","of","the","output","of","model","'s","stem",".","representation_size",":","Size","of","the","representation","layer","in","the","model","'s","head",".","if","None",",","we","skip","the","extra","projection","+","tanh","activation","at","the","end",".","temporal_encoding_config",":","ConfigDict","which","defines","the","type","of","input","encoding","when","tokenising","the","video",".","attention_config",":","ConfigDict","which","defines","the","type","of","spatio-temporal","attention","applied","in","the","model",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","stochastic_droplayer_rate",":","Probability","of","dropping","a","layer",".","Linearly","increases","from","0","to","the","provided","value","..","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",".","return_prelogits",":","If","true",",","return","the","final","representation","of","the","network","before","the","classification","head",".","Useful","when","using","features","for","a","downstream","task",".","return_preclassifier",":","If","true",",","return","a","dict","of","all","token","embeddings",".","Useful","when","using","token","embeddings","for","a","downstream","task",".","return_as_dict",":","If","true",",","return","the","token","embeddings","as","a","dictionary","instead","of","a","concatenated","tensor",".","use_bottleneck",":","If","True",",","adds","self-attention","bottleneck",".","n_bottlenecks",":","Number","of","bottleneck","tokens",".","test_with_bottlenecks",":","Whether","to","use","bottlenecks","at","test","time",".","share_encoder",":","If","True",",","different","modalities","share","the","same","encoder","weights","for","the","layers","before","fusion",".","return_bottlenecks",":","If","True",",","return","bottleneck","embeddings",".","use_modality_tokens",":","If","True",",","modality","tokens","are","used",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self,
      x: Dict[str, jnp.ndarray],
      *,
      train: bool,
      debug: bool = False,
  ):
      assert (
          self.fusion_layer <= self.num_layers
          and self.fusion_layer >= 0
      )
      assert self.classifier in [
          "onetoken",
          "token",
          "0",
          "gap",
          "gmp",
          "gsp",
      ]
      attention_type = self.attention_config.get("type", "spacetime")
      assert attention_type not in [
          "factorized_transformer_block",
          "factorized_self_attention_block",
          "factorized_dot_product_attention",
      ], "Factorised attention is not implemented"
  
      x, temporal_dims = self._temporal_encode(x)
      if self.use_modality_tokens:
          x = self.add_modality_token(x)
  
      bottleneck_dtype = x[self.modality_fusion[0]].dtype
      n, _, c = x[self.modality_fusion[0]].shape
      bottleneck = None
      if self.use_bottleneck:
          n_bottlenecks = self.n_bottlenecks
          if self.classifier in ["token"]:
              n_bottlenecks += 1
          bottleneck = self.param(
              "bottleneck",
              nn.initializers.normal(stddev=0.02),  # From BERT.
              (1, n_bottlenecks, c),
              bottleneck_dtype,
          )
          bottleneck = jnp.tile(bottleneck, [n, 1, 1])
  
      token_lengths = {m: x[m].shape[1] for m in self.modality_fusion}
      output = Encoder(
          modality_fusion=self.modality_fusion,
          fusion_layer=self.fusion_layer,
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          attention_config=self.attention_config,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_droplayer_rate=self.stochastic_droplayer_rate,
          use_bottleneck=self.use_bottleneck,
          test_with_bottlenecks=self.test_with_bottlenecks,
          share_encoder=self.share_encoder,
          return_bottlenecks=self.return_bottlenecks,
          dtype=self.dtype,
          name="Transformer",
      )(x, bottleneck, train=train)
      if self.return_bottlenecks:
          x, bottleneck = output
      else:
          x = output
      if self.return_preclassifier:
          if self.return_as_dict:
              x_dict = {}
              for m in self.modality_fusion:
                  v = token_lengths[m]
                  x_dict[m] = x[:, :v]
                  x = x[:, v:]
              x = x_dict
          if self.return_bottlenecks:
              return x, bottleneck
          return x
  
      x_out = ProjectionHeadAndClassifier(
          classifier=self.classifier,
          num_classes=self.num_classes,
          modality_fusion=self.modality_fusion,
          representation_size=self.representation_size,
          return_prelogits=self.return_prelogits,
      )(x, temporal_dims, train=train)
      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 3 for calling the MBT obj's underlying member method code with example usage: obj.__call__(x) and returns: x_out, x, x,  bottleneck