class_docstr: Axial Self Attention module.
  
  Attributes:
    attention_axis: Axis of the attention
    axial_attention_configs: Configurations of the axial attention.
class_docstr_tok: ["Axial","Self","Attention","module",".","Attributes",":","attention_axis",":","Axis","of","the","attention","axial_attention_configs",":","Configurations","of","the","axial","attention","."]
docstr: Applies Axial Self Attention module.
  
  Args:
    x: Input data.
  
  Returns:
    Output after axial attention applied on it.
docstr_tok: ["Applies","Axial","Self","Attention","module",".","Args",":","x",":","Input","data",".","Returns",":","Output","after","axial","attention","applied","on","it","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      """Applies Axial Self Attention module.
  
      Args:
        x: Input data.
  
      Returns:
        Output after axial attention applied on it.
      """
      bs, height, width, channel = x.shape
      if self.attention_axis == 1:  # Row attention.
          x = x.transpose((0, 2, 1, 3)).reshape(
              bs * width, height, channel
          )
      elif self.attention_axis == 2:  # Column attention.
          x = x.reshape(bs * height, width, channel)
      else:
          raise ValueError(
              "Only attention over rows or columns is supported."
          )
  
      x = SelfAttentionWith1DRelativePos(
          self.axial_attention_configs.num_heads
      )(x)
  
      if self.attention_axis == 1:
          return x.reshape((bs, width, height, channel)).transpose(
              (0, 2, 1, 3)
          )
      else:
          return x.reshape((bs, height, width, channel))
  
imports: from typing import Dict, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.model_lib.layers import nn_ops
  from scenic.projects.baselines import bit_resnet
label_desc: Use image node_id 1 for calling the AxialSelfAttention obj's underlying member method code with example usage: obj.__call__(x) and returns: x, x