class_docstr: Transformer Encoder.
  
  Attributes:
    inputs: nd-array, Input data
    modality_fusion: Tuple with modalities to combine.
    mlp_dim: Dimension of the mlp on top of attention block.
    num_layers: Number of layers.
    num_heads: Number of attention heads.
    attention_config: Has parameters for the type of attention.
    dropout_rate: Dropout rate.
    fusion_layer: Which layer to fuse modalities. fusion_layer == 0 provides
      early fusion.
    attention_dropout_rate: Dropout for attention heads.
    stochastic_droplayer_rate: Probability of dropping a layer linearly
      grows from 0 to the provided value. Our implementation of stochastic
      depth follows timm library, which does per-example layer dropping and
      uses independent dropping patterns for each skip-connection.
    use_bottleneck: If True, adds self-attention bottleneck.
    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.
    add_pos_embedding: If True, positional embeddings are added to the input
      token embeddings.
    return_bottlenecks: If True, return bottleneck embeddings.
class_docstr_tok: ["Transformer","Encoder",".","Attributes",":","inputs",":","nd-array",",","Input","data","modality_fusion",":","Tuple","with","modalities","to","combine",".","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_layers",":","Number","of","layers",".","num_heads",":","Number","of","attention","heads",".","attention_config",":","Has","parameters","for","the","type","of","attention",".","dropout_rate",":","Dropout","rate",".","fusion_layer",":","Which","layer","to","fuse","modalities",".","fusion_layer","==","0","provides","early","fusion",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","stochastic_droplayer_rate",":","Probability","of","dropping","a","layer","linearly","grows","from","0","to","the","provided","value",".","Our","implementation","of","stochastic","depth","follows","timm","library",",","which","does","per-example","layer","dropping","and","uses","independent","dropping","patterns","for","each","skip-connection",".","use_bottleneck",":","If","True",",","adds","self-attention","bottleneck",".","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",".","add_pos_embedding",":","If","True",",","positional","embeddings","are","added","to","the","input","token","embeddings",".","return_bottlenecks",":","If","True",",","return","bottleneck","embeddings","."]
docstr: Applies Transformer model on the inputs.
code: def __call__(
      self, x: Dict[str, Any], bottleneck: jnp.ndarray, *, train: bool
  ):
      """Applies Transformer model on the inputs."""
  
      def get_encoder_block(encoder_block, droplayer_p, name):
          """Returns the encoder block for a single layer."""
          dtype = jax.dtypes.canonicalize_dtype(self.dtype)
          return encoder_block(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              droplayer_p=droplayer_p,
              name=name,
              dtype=dtype,
          )
  
      def get_context(target_modality, modality_fusion, x):
          """Returns list of context modalities."""
          context = []
          for modality in _MODALITIES:
              if (
                  modality != target_modality
                  and modality in modality_fusion
              ):
                  context.append(x[modality])
          return context
  
      def combine_context(x, other_modalities):
          """Combine x with a list of other modalities."""
          t_x = x.shape[1]
          # Append x to the end of the list
          other_modalities.append(x)
          x_combined = jnp.concatenate(other_modalities, axis=1)
          return x_combined, t_x
  
      assert self.modality_fusion
  
      # Add positional embeddings
      if self.add_pos_embedding:
          for modality in self.modality_fusion:
              if modality == "spectrogram":
                  modality_name = "spec"
              else:
                  modality_name = modality
              if modality == "rgb":
                  name = ""
              else:
                  name = "_" + modality_name
              x[modality] = add_positional_embed(
                  x[modality], "posembed_input" + name
              )
  
      if (
          self.attention_config is None
          or self.attention_config.type
          in [  # pytype: disable=attribute-error
              "spacetime",
              "factorized_encoder",
          ]
      ):
          encoder_block = mbt_model.EncoderBlock
      else:
          raise ValueError(
              f"Unknown attention type {self.attention_config.type}"
          )  # pytype: disable=attribute-error
  
      use_bottlenecks = train or self.test_with_bottlenecks
      x_combined = None
      # Input Encoder
      for lyr in range(self.num_layers):
          droplayer_p = (
              lyr / max(self.num_layers - 1, 1)
          ) * self.stochastic_droplayer_rate
          encoders = {}
          first_modality = self.modality_fusion[0]
          encoders[first_modality] = get_encoder_block(
              encoder_block, droplayer_p, f"encoderblock_{lyr}"
          )
          for modality in self.modality_fusion:
              # This is important for loading old checkpoints, where we used 'spec'
              if modality == "spectrogram":
                  modality_name = "spec"
              else:
                  modality_name = modality
              if modality != first_modality:
                  if self.share_encoder:
                      encoders[modality] = encoders[first_modality]
                  else:
                      encoders[modality] = get_encoder_block(
                          encoder_block,
                          droplayer_p,
                          f"encoderblock_{lyr}_" + modality_name,
                      )
  
          if (
              lyr < self.fusion_layer
              or len(self.modality_fusion) == 1
              or (self.use_bottleneck and not use_bottlenecks)
          ):
              for modality in self.modality_fusion:
                  x[modality] = encoders[modality](
                      x[modality], deterministic=not train
                  )
          else:
              if self.use_bottleneck:
                  bottle = []
                  for modality in self.modality_fusion:
                      t_mod = x[modality].shape[1]
                      in_mod = jnp.concatenate(
                          [x[modality], bottleneck], axis=1
                      )
                      out_mod = encoders[modality](
                          in_mod, deterministic=not train
                      )
                      x[modality] = out_mod[:, :t_mod]
                      bottle.append(out_mod[:, t_mod:])
                  bottleneck = jnp.mean(
                      jnp.stack(bottle, axis=-1), axis=-1
                  )
              else:
                  if (
                      not self.share_encoder
                      and len(self.modality_fusion) > 1
                  ):
                      x_new = {}
                      for modality in self.modality_fusion:
                          other_modalities = get_context(
                              modality, self.modality_fusion, x
                          )
                          combined_mods, t = combine_context(
                              x[modality], other_modalities
                          )
                          combined_mods = encoders[modality](
                              combined_mods, deterministic=not train
                          )
                          x_new[modality] = combined_mods[:, -t:]
                      x = x_new
  
                  elif (
                      self.share_encoder
                      and len(self.modality_fusion) > 1
                  ):
                      if x_combined is None:
                          x_combined = []
                          for modality in self.modality_fusion:
                              x_combined.append(x[modality])
                          x_combined = jnp.concatenate(
                              x_combined, axis=1
                          )
                      x_combined = encoders[first_modality](
                          x_combined, deterministic=not train
                      )
      if x_combined is not None:
          x_out = x_combined
      else:
          x_out = []
          for modality in self.modality_fusion:
              x_out.append(x[modality])
          x_out = jnp.concatenate(x_out, axis=1)
      encoded = nn.LayerNorm(name="encoder_norm")(x_out)
  
      if self.return_bottlenecks:
          assert self.use_bottleneck, (
              "`use_bottleneck' should be True to return " "bottlenecks"
          )
          return encoded, bottleneck
  
      return encoded
  
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 Encoder obj's underlying member method code with example usage: obj.__call__(x, bottleneck) and returns: encoded, encoder_block, context, x_combined,  t_x, encoded,  bottleneck