class_docstr: 1-Dimensional Transformer encoder block.
  
  Attributes:
    mlp_dim: dimension of the MLP on top of attention block.
    attention_configs: Configs pass to the self-attention func.
    attention_fn: Type of the self-attention function.
    dropout_rate: Dropout used in the MLP block.
    attention_dropout_rate: Dropout for attention heads.
    num_kernel_features: Number of kernel features used.
    redraw: Whether to redraw (valid only if random faturees are used).
    post_sa_fn: Function to be applied on the output of self-attention block.
    dtype: The dtype of the computation.
class_docstr_tok: ["1-Dimensional","Transformer","encoder","block",".","Attributes",":","mlp_dim",":","dimension","of","the","MLP","on","top","of","attention","block",".","attention_configs",":","Configs","pass","to","the","self-attention","func",".","attention_fn",":","Type","of","the","self-attention","function",".","dropout_rate",":","Dropout","used","in","the","MLP","block",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","num_kernel_features",":","Number","of","kernel","features","used",".","redraw",":","Whether","to","redraw","(","valid","only","if","random","faturees","are","used",")",".","post_sa_fn",":","Function","to","be","applied","on","the","output","of","self-attention","block",".","dtype",":","The","dtype","of","the","computation","."]
docstr: Applies Encoder1DBlock module.
  
  Args:
    inputs_q: Input data in shape of `[bs, len, c]`.
    inputs_kv: Memory data in shape of `[bs, memory len, c]`.
    deterministic: Whether the model is in deterministic mode (if so, do not
      apply dropout).
  
  Returns:
    Output after 1-d transformer encoder block.
code: def __call__(
      self,
      inputs_q: jnp.ndarray,
      inputs_kv: jnp.ndarray,
      *,
      deterministic: bool,
  ) -> jnp.ndarray:
      """Applies Encoder1DBlock module.
  
      Args:
        inputs_q: Input data in shape of `[bs, len, c]`.
        inputs_kv: Memory data in shape of `[bs, memory len, c]`.
        deterministic: Whether the model is in deterministic mode (if so, do not
          apply dropout).
  
      Returns:
        Output after 1-d transformer encoder block.
      """
      assert inputs_q.ndim == 3
      if self.attention_fn:
          is_self_attention = inputs_kv is None
  
          # Attention block.
          valid_args = _get_variant_args(self.attention_fn)
          # Remove args that are potentially not needed for variant.
          attention_configs = {
              x: self.attention_configs[x]
              for x in valid_args
              if x in self.attention_configs
          }
          x = nn.LayerNorm(dtype=self.dtype)(inputs_q)
          if not is_self_attention:
              assert inputs_kv.ndim == 3
              inputs_kv = nn.LayerNorm(dtype=self.dtype)(inputs_kv)
  
          # Prepare the input for the attention modole.
          # We shouldn't pass memory if it is self-attention.
          init_arg_to_attention_module = {
              "kernel_init": nn.initializers.xavier_uniform(),
              "broadcast_dropout": False,
              "dtype": self.dtype,
              "dropout_rate": self.attention_dropout_rate,
          }
          inputs_to_attention_module = {
              "inputs_q": x,
              "deterministic": deterministic,
          }
          if is_self_attention:
              inputs_to_attention_module["inputs_kv"] = x
          else:
              inputs_to_attention_module["inputs_kv"] = inputs_kv
  
          x = _get_attention_module(
              self.attention_fn, is_self_attention
          )(**init_arg_to_attention_module, **attention_configs)(
              **inputs_to_attention_module
          )
  
          x = nn.Dropout(rate=self.dropout_rate)(
              x, deterministic=deterministic
          )
  
          if x.shape[-2] != inputs_q.shape[-2]:
              # TODO(yitay): Support case where we downsample. How do we handle res?
              # Currently bypassing this causes training problems...
              raise ValueError(
                  "Shape not identical. Cannot add residual connection."
              )
  
          drop_pattern = self.get_drop_pattern(x, deterministic)
          x = x * (1.0 - drop_pattern) + inputs_q
          if self.post_sa_fn is not None:
              x = self.post_sa_fn(x)  # pylint: disable=not-callable
      else:
          x = inputs_q
  
      if self.mlp_dim is None:
          # Skip the MLP block.
          return x
  
      # MLP block.
      y = nn.LayerNorm(dtype=self.dtype)(x)
      mlp_dim = self.mlp_dim
      if isinstance(self.mlp_dim, int):
          mlp_dim = (mlp_dim,)
      for mlp_d in mlp_dim:
          y = attention_layers.MlpBlock(
              mlp_dim=mlp_d,
              dtype=self.dtype,
              dropout_rate=self.dropout_rate,
              activation_fn=nn.gelu,
              kernel_init=nn.initializers.xavier_uniform(),
              bias_init=nn.initializers.normal(stddev=1e-6),
          )(y, deterministic=deterministic)
  
      drop_pattern = self.get_drop_pattern(x, deterministic)
      return x + y * (1.0 - drop_pattern)
  
imports: import abc
  import enum
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Iterable
  from absl import logging
  from flax import linen
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 2 for calling the Encoder1DBlock obj's underlying member method code with example usage: obj.__call__(inputs_q, inputs_kv) and returns: unknown, x