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.
class_docstr_tok: ["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","."]
docstr: Applies Transformer model on the inputs.
  
  Args:
    inputs: Input tokens of shape [batch, num_tokens, channels].
    matvit_mask_dims: matvit nesting dimensions, a list of size num_layers.
    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","]",".","matvit_mask_dims",":","matvit","nesting","dimensions",",","a","list","of","size","num_layers",".","train",":","If","in","training","mode",",","dropout","and","stochastic","depth","is","applied",".","Returns",":","Encoded","tokens","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      *,
      matvit_mask_dims: List[float],
      train: bool = False,
  ):
      """Applies Transformer model on the inputs.
  
      Args:
        inputs: Input tokens of shape [batch, num_tokens, channels].
        matvit_mask_dims: matvit nesting dimensions, a list of size num_layers.
        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 = 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_2d":
          batch, num_tokens, hidden_dim = inputs.shape
          height = width = int(np.sqrt(num_tokens))
          if height * width != num_tokens:
              raise ValueError(
                  "Input is assumed to be square for sinusoidal init."
              )
          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):
          nesting_mask = jnp.ones(
              matvit_mask_dims[lyr], dtype=jnp.int32
          )
          ffn_mask = jnp.zeros(self.mlp_dim, dtype=jnp.int32)
          ffn_mask = jax.lax.dynamic_update_slice(
              ffn_mask, nesting_mask, (0,)
          )
          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, matvit_mask=ffn_mask)
      encoded = nn.LayerNorm(name="encoder_norm")(x)
      return encoded
  
imports: from typing import Any, Callable, List, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  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.matvit import layers
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