class_docstr: Audio-Visual Fusion Transformer model for Video.
  
  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_layers: Number of layers.
    num_heads: Number of self-attention heads.
    num_classes: Number of output classes.
    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.
      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.
    representation_size: Size of the representation layer in the model's head.
    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'.
    modality_fusion: Tuple with modalities to combine.
    fusion_layer: Which layer to fuse modalities.
    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.
    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.
    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_layers",":","Number","of","layers",".","num_heads",":","Number","of","self-attention","heads",".","num_classes",":","Number","of","output","classes",".","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",".","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",".","representation_size",":","Size","of","the","representation","layer","in","the","model","'s","head",".","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","'",".","modality_fusion",":","Tuple","with","modalities","to","combine",".","fusion_layer",":","Which","layer","to","fuse","modalities",".","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",".","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",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(self, x, *, train: bool, debug: bool = False):
      assert (
          self.fusion_layer <= self.num_layers
          and self.fusion_layer >= 0
      )
      assert self.classifier in ["token", "0", "gap", "gmp", "gsp"]
  
      temporal_dims = {}
      for modality in self.modality_fusion:
          x[modality], _ = temporal_encode(
              x[modality],
              modality,
              self.temporal_encoding_config,
              self.patches,
              self.hidden_size,
          )
          # If we want to add a class token, add it here.
          if self.classifier in ["token"]:
              if modality == "rgb" or len(self.modality_fusion) == 1:
                  name = ""
              else:
                  name = modality
              n, temporal_dims[modality], c = x[modality].shape
              cls = self.param(
                  "cls" + name,
                  nn.initializers.zeros,
                  (1, 1, c),
                  x[modality].dtype,
              )
              cls = jnp.tile(cls, [n, 1, 1])
              x[modality] = jnp.concatenate([cls, x[modality]], axis=1)
              bottleneck_dtype = x[modality].dtype
  
      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])
  
      x = 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,
          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,
          dtype=self.dtype,
          name="Transformer",
      )(x, bottleneck, train=train)
  
      if self.return_preclassifier:
          return x
  
      if self.classifier in ["token", "0"]:
          # Obtaining the CLS tokens for each modality.
          x_out = {}
          counter = 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
          ]
          x_out = fn(x, axis=list(range(1, x.ndim - 1)))
  
      if self.representation_size is not None:
          pre_logits_fc = nn.Dense(
              self.representation_size, name="pre_logits"
          )
          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 = nn.Dense(
                  self.representation_size, name="pre_logits"
              )(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)
          if not train:
              return x_pool
      else:
          x_out = nn.Dense(
              self.num_classes,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )(x_out)
      return x_out
  
imports: import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  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 attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.mbt import model_utils
label_desc: Use image node_id 1 for calling the MBT obj's underlying member method code with example usage: obj.__call__(x) and returns: x_out, x, x_out, x_pool