class_docstr: Transformer Encoder.
  
  Attributes:
    inputs: nd-array, Input data
    temporal_dims: Number of temporal dimensions in the input.
    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.
    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.
    positional_embedding: The type of positional embedding to use. Supported
      values are {learned_1d, sinusoidal_1d, sinusoidal_3d, none}.
    normalise_output: If True, perform layernorm on the output.
class_docstr_tok: ["Transformer","Encoder",".","Attributes",":","inputs",":","nd-array",",","Input","data","temporal_dims",":","Number","of","temporal","dimensions","in","the","input",".","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",".","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",".","positional_embedding",":","The","type","of","positional","embedding","to","use",".","Supported","values","are","{","learned_1d",",","sinusoidal_1d",",","sinusoidal_3d",",","none","}",".","normalise_output",":","If","True",",","perform","layernorm","on","the","output","."]
docstr: Applies Transformer model on the inputs.
code: def __call__(self, inputs: jnp.ndarray, *, train: bool):
      """Applies Transformer model on the inputs."""
      assert inputs.ndim == 3  # (batch, len, emb)
      dtype = jax.dtypes.canonicalize_dtype(self.dtype)
  
      if self.positional_embedding == "learned_1d":
          x = vit.AddPositionEmbs(
              posemb_init=nn.initializers.normal(
                  stddev=0.02
              ),  # from BERT.
              name="posembed_input",
          )(inputs)
      elif self.positional_embedding == "sinusoidal_1d":
          x = attention_layers.Add1DPositionEmbedding(posemb_init=None)(
              inputs
          )
      elif self.positional_embedding == "sinusoidal_3d":
          batch, num_tokens, hidden_dim = inputs.shape
          height = width = int(
              np.sqrt(num_tokens // self.temporal_dims)
          )
          if height * width * self.temporal_dims != num_tokens:
              raise ValueError(
                  "Input is assumed to be square for sinusoidal init."
              )
          inputs_reshape = inputs.reshape(
              [batch, self.temporal_dims, height, width, hidden_dim]
          )
          x = attention_layers.AddFixedSinCosPositionEmbedding()(
              inputs_reshape
          )
          x = x.reshape([batch, num_tokens, hidden_dim])
      elif self.positional_embedding == "none":
          x = inputs
      else:
          raise ValueError(
              f"Unknown positional embedding {self.positional_embedding}"
          )
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      if (
          self.attention_config is None
          or self.attention_config.type
          in ["spacetime", "factorized_encoder"]
      ):
          encoder_block = EncoderBlock
      elif (
          self.attention_config.type
          == "factorized_self_attention_block"
      ):
          encoder_block = functools.partial(
              EncoderFactorizedSelfAttentionBlock,
              attention_order=self.attention_config.attention_order,
              attention_kernel_initializer=KERNEL_INITIALIZERS[
                  self.attention_config.get(
                      "attention_kernel_init_method", "xavier"
                  )
              ],
              temporal_dims=self.temporal_dims,
          )
      elif (
          self.attention_config.type
          == "factorized_dot_product_attention"
      ):
          b, thw, d = x.shape
          x = _reshape_to_time_space(
              x, self.temporal_dims
          )  # [b, t, hw, d]
          encoder_block = functools.partial(
              EncoderBlock,
              attention_fn=functools.partial(
                  model_utils.factorized_dot_product_attention
              ),
          )
      else:
          raise ValueError(
              f"Unknown attention type {self.attention_config.type}"
          )
  
      # Input Encoder
      for lyr in range(self.num_layers):
          droplayer_p = (
              lyr / max(self.num_layers - 1, 1)
          ) * self.stochastic_droplayer_rate
          x = 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=f"encoderblock_{lyr}",
              dtype=dtype,
          )(x, deterministic=not train)
  
      if (
          self.attention_config.type
          == "factorized_dot_product_attention"
      ):
          # Reshape back to 3D:
          x = jnp.reshape(x, [b, thw, d])
  
      if self.normalise_output:
          encoded = nn.LayerNorm(name="encoder_norm")(x)
      else:
          encoded = x
  
      return encoded
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  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.common_lib import video_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 attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 1 for calling the Encoder obj's underlying member method code with example usage: obj.__call__(inputs) and returns: encoded