class_docstr: Transformer Axial Encoder.
  
  EncoderAxial replaces each Transformer block, i.e. `self-attention + MLP` with
  two axial attention block: `[self-row-attention + MLP] + [self-col-attention
  + MLP]`.
  Attributes:
    mlp_dim: Dimension of the MLP on top of attention block.
    num_layers: Number of layers.
    transformer_encoder_type: Type of the transformer encoder, one of
      ['axial', 'axial_attention'].
    attention_configs: Configurations passed to the self-attention.
    attention_fn: Self-attention function used in the model.
    dropout_rate:  Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
class_docstr_tok: ["Transformer","Axial","Encoder",".","EncoderAxial","replaces","each","Transformer","block",",","i.e",".","`","self-attention","+","MLP","`","with","two","axial","attention","block",":","`","[","self-row-attention","+","MLP","]","+","[","self-col-attention","+","MLP","]","`",".","Attributes",":","mlp_dim",":","Dimension","of","the","MLP","on","top","of","attention","block",".","num_layers",":","Number","of","layers",".","transformer_encoder_type",":","Type","of","the","transformer","encoder",",","one","of","[","'axial","'",",","'axial_attention","'","]",".","attention_configs",":","Configurations","passed","to","the","self-attention",".","attention_fn",":","Self-attention","function","used","in","the","model",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads","."]
docstr: Applies Axial Transformer model on the inputs.
  
  Args:
    inputs: Input data.
    train: If it is training.
  
  Returns:
    Output of a transformer encoder.
docstr_tok: ["Applies","Axial","Transformer","model","on","the","inputs",".","Args",":","inputs",":","Input","data",".","train",":","If","it","is","training",".","Returns",":","Output","of","a","transformer","encoder","."]
code: def __call__(
      self, inputs: jnp.ndarray, train: bool = False
  ) -> jnp.ndarray:
      """Applies Axial Transformer model on the inputs.
  
      Args:
        inputs: Input data.
        train: If it is training.
  
      Returns:
        Output of a transformer encoder.
      """
      assert inputs.ndim == 4  # Shape is `[batch, h, w, emb]`.
      x = attention_layers.Add2DPositionEmbedding(
          posemb_init=nn.initializers.normal(stddev=0.02),  # From BERT.
          name="posembed_input",
      )(inputs)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      # Input Encoder
      for lyr in range(self.num_layers):
          if self.transformer_encoder_type == "axial":
              two_d_shape = x.shape
              # Row attention.
              x = model_utils.get_axial_1d_input(x, axis=1)
              x = model_utils.Encoder1DBlock(  # pytype: disable=wrong-arg-types  # jax-ndarray
                  mlp_dim=self.mlp_dim,
                  attention_fn=self.attention_fn,
                  attention_configs=self.attention_configs,
                  dropout_rate=self.dropout_rate,
                  attention_dropout_rate=self.attention_dropout_rate,
                  name=f"encoder_block_row_{lyr}",
              )(
                  x, inputs_kv=None, deterministic=not train
              )
              x = model_utils.get_axial_2d_input(
                  x, axis=1, two_d_shape=two_d_shape
              )
  
              # Column attention.
              x = model_utils.get_axial_1d_input(x, axis=2)
              x = model_utils.Encoder1DBlock(  # pytype: disable=wrong-arg-types  # jax-ndarray
                  mlp_dim=self.mlp_dim,
                  attention_fn=self.attention_fn,
                  attention_configs=self.attention_configs,
                  dropout_rate=self.dropout_rate,
                  attention_dropout_rate=self.attention_dropout_rate,
                  name=f"encoder_block_col_{lyr}",
              )(
                  x, inputs_kv=None, deterministic=not train
              )
              x = model_utils.get_axial_2d_input(
                  x, axis=2, two_d_shape=two_d_shape
              )
  
          elif self.transformer_encoder_type == "axial_attention":
              x = model_utils.EncoderAxialBlock(
                  mlp_dim=self.mlp_dim,
                  attention_fn=self.attention_fn,
                  attention_configs=self.attention_configs,
                  dropout_rate=self.dropout_rate,
                  attention_dropout_rate=self.attention_dropout_rate,
                  name=f"axial_encoder_fblock_{lyr}",
              )(x, deterministic=not train)
  
          else:
              raise ValueError(
                  "Undefined transformer encoder type: "
                  f"{self.transformer_encoder_type}."
              )
  
      return nn.LayerNorm(x, name="encoder_norm")
  
imports: from typing import Any, Optional
  import flax.linen
  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
  from scenic.model_lib.layers import nn_ops
  from scenic.projects.baselines import vit
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 1 for calling the EncoderAxial obj's underlying member method code with example usage: obj.__call__(inputs, train) and returns: nn