class_docstr: Transformer Encoder.
  
  Attributes:
    num_layers: Number of layers.
    mlp_dim: Dimension of the mlp on top of attention block.
    num_heads: The number of heads for multi-head self-attention.
    positional_embedding: The type of positional embeddings to add to the
      input tokens. Options are {learned_1d, sinusoidal_2d, none}.
    dropout_rate: Dropout rate.
    stochastic_depth: 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.
    dtype: Dtype of activations.
    has_cls_token: Whether or not the sequence is prepended with a CLS token.
docstr: Applies Transformer model on the inputs.
  
  Args:
    inputs: Input tokens of shape [batch, num_tokens, channels].
    train: If in training mode, dropout and stochastic depth is applied.
  
  Returns:
    Encoded tokens.
docstr_tok: ["Applies","Transformer","model","on","the","inputs",".","Args",":","inputs",":","Input","tokens","of","shape","[","batch",",","num_tokens",",","channels","]",".","train",":","If","in","training","mode",",","dropout","and","stochastic","depth","is","applied",".","Returns",":","Encoded","tokens","."]
code: def __call__(self, inputs: jnp.ndarray, *, train: bool = False):
      """Applies Transformer model on the inputs.
  
      Args:
        inputs: Input tokens of shape [batch, num_tokens, channels].
        train: If in training mode, dropout and stochastic depth is applied.
  
      Returns:
        Encoded tokens.
      """
  
      assert inputs.ndim == 3  # Shape is `[batch, len, emb]`.
      dtype = jax.dtypes.canonicalize_dtype(self.dtype)
  
      # Add positional embeddings to tokens.
      if self.positional_embedding == "learned_1d":
          x = 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_2d":
          batch, num_tokens, hidden_dim = inputs.shape
          if self.has_cls_token:
              num_tokens -= 1
          height = width = int(np.sqrt(num_tokens))
          if height * width != num_tokens:
              raise ValueError(
                  "Input is assumed to be square for sinusoidal init."
              )
          if self.has_cls_token:
              inputs_reshape = inputs[:, 1:].reshape(
                  [batch, height, width, hidden_dim]
              )
              x = attention_layers.AddFixedSinCosPositionEmbedding()(
                  inputs_reshape
              )
              x = x.reshape([batch, num_tokens, hidden_dim])
              x = jnp.concatenate([inputs[:, :1], x], axis=1)
          else:
              inputs_reshape = inputs.reshape(
                  [batch, 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(
              "Unknown positional embedding: "
              f"{self.positional_embedding}"
          )
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      # Input Encoder.
      for lyr in range(self.num_layers):
          x = Encoder1DBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              stochastic_depth=(lyr / max(self.num_layers - 1, 1))
              * self.stochastic_depth,
              name=f"encoderblock_{lyr}",
              dtype=dtype,
          )(x, deterministic=not train)
      encoded = nn.LayerNorm(name="encoder_norm")(x)
      return encoded
  
imports: from typing import Any, Callable, Optional, Sequence
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  import scipy
  from tensorflow.io import gfile
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