class_docstr: Transformer Encoder.
  
  Attributes:
    num_layers: Number of layers.
    mlp_dim: Dimension of the mlp on top of attention block.
    num_heads: Number of self-attention heads.
    ac_config: Configuration of the adaptive computation.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Attention 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.
docstr_tok: ["Applies","Transformer","model","on","the","inputs","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      bank: Optional[jnp.ndarray] = None,
      *,
      input_mask: Optional[jnp.ndarray] = None,
      train: bool = False,
  ):
      """Applies Transformer model on the inputs."""
  
      assert x.ndim == 3  # Shape is `[batch, len, emb]`.
      # Init it as None and update it later.
      logging_input_mask = None
      tape_added = 0
      loss_atr = None
      # Input Encoder.
      for lyr in range(self.num_layers):
          if self.ac_config.get("add_tape_token_to_layers", []):
              output_q_kv = get_q_kv_mask(
                  x,
                  input_mask,
                  lyr,
                  tape_added,
                  self.ac_config,
                  bank,
                  train=train,
              )
              (
                  x_q,
                  x_kv,
                  input_mask,
                  logging_input_mask_tmp,
                  loss_atr_tmp,
                  tape_added,
              ) = output_q_kv
              # Update logging_input_mask only when the returned value is not None.
              if logging_input_mask_tmp is not None:
                  logging_input_mask = logging_input_mask_tmp
              # Update loss_atr only when the returned value is not None.
              if loss_atr_tmp is not None:
                  loss_atr = loss_atr_tmp
          else:
              # Add no tape token
              x_q, x_kv = x, None
  
          x = Encoder1DBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              ac_config=self.ac_config,
              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=jax.dtypes.canonicalize_dtype(self.dtype),
          )(
              x_q,
              x_kv,
              input_mask=input_mask,
              added_tape_len=tape_added,
              deterministic=not train,
          )
  
      encoded = nn.LayerNorm(name="encoder_norm")(x)
      return encoded, (logging_input_mask, loss_atr)
  
imports: from typing import Any, Callable, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the AdaTapeEncoder obj's underlying member method code with example usage: obj.__call__(x, bank) and returns: encoded,  logging_input_mask,  loss_atr