class_docstr: 2-Dimensional Transformer encoder block with Axial attention.
  
  This block is similar to Encoder1DBlock, where instead of `self-attention+MLP`
  we have `row-self-attention + col-self-attention + MLP`.
  
  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 sel-attention function.
    dropout_rate: Dropout used in the mlp block.
    attention_dropout_rate: Dropout for attention heads.
    factorization_axis: Axis over which we run attention.
    post_sa_fn: Function to be applied on the output of self-attention block.
    dtype: The dtype of the computation.
class_docstr_tok: ["2-Dimensional","Transformer","encoder","block","with","Axial","attention",".","This","block","is","similar","to","Encoder1DBlock",",","where","instead","of","`","self-attention+MLP","`","we","have","`","row-self-attention","+","col-self-attention","+","MLP","`",".","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","sel-attention","function",".","dropout_rate",":","Dropout","used","in","the","mlp","block",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","factorization_axis",":","Axis","over","which","we","run","attention",".","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: Input data in shape of `[bs, len, c]`.
    deterministic: Whether the model is in deterministic mode (if so, do not
      apply dropout).
  
  Returns:
    Output after axial attention encoder block.
code: def __call__(
      self, inputs: jnp.ndarray, *, deterministic: bool
  ) -> jnp.ndarray:
      """Applies Encoder1DBlock module.
  
      Args:
        inputs: Input data in shape of `[bs, len, c]`.
        deterministic: Whether the model is in deterministic mode (if so, do not
          apply dropout).
  
      Returns:
        Output after axial attention encoder block.
      """
  
      def _run_attention_on_axis(inputs, axis, two_d_shape):
          inputs = get_axial_1d_input(inputs, axis=axis)
          x = nn.LayerNorm(dtype=self.dtype)(inputs)
          init_arg_to_attention_module = {
              "kernel_init": nn.initializers.xavier_uniform(),
              "broadcast_dropout": False,
              "dtype": self.dtype,
              "dropout_rate": self.attention_dropout_rate,
          }
          # 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 = _get_attention_module(
              self.attention_fn, is_self_attention=True
          )(**init_arg_to_attention_module, **attention_configs)(
              x, deterministic=deterministic
          )
          x = nn.Dropout(rate=self.dropout_rate)(
              x, deterministic=deterministic
          )
          drop_pattern = self.get_drop_pattern(x, deterministic)
          x = x * (1.0 - drop_pattern) + inputs
          return get_axial_2d_input(
              x, axis=axis, two_d_shape=two_d_shape
          )
  
      x = inputs
      if self.attention_fn:
          # Row attention block.
          two_d_shape = inputs.shape
  
          for ax in self.factorization_axis:
              x = _run_attention_on_axis(x, ax, two_d_shape)
  
          if self.post_sa_fn is not None:
              x = self.post_sa_fn(x)  # pylint: disable=not-callable
  
      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 EncoderAxialBlock obj's underlying member method code with example usage: obj.__call__(inputs) and returns: unknown, get_axial_2d_input, x