class_docstr: Bottleneck AxialResNet block.
  
  Attributes:
    nout: Number of output features.
    axial_attention_configs: Configurations of the axial attention.
    strides: Down-sampling stride.
    bottleneck: If True, the block is a bottleneck block.
class_docstr_tok: ["Bottleneck","AxialResNet","block",".","Attributes",":","nout",":","Number","of","output","features",".","axial_attention_configs",":","Configurations","of","the","axial","attention",".","strides",":","Down-sampling","stride",".","bottleneck",":","If","True",",","the","block","is","a","bottleneck","block","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      features = self.nout
      nout = self.nout * 4 if self.bottleneck else self.nout
      needs_projection = x.shape[-1] != nout or self.strides != (1, 1)
      residual = x
      if needs_projection:
          residual = bit_resnet.StdConv(
              nout,
              (1, 1),
              self.strides,
              use_bias=False,
              name="conv_proj",
          )(residual)
          residual = nn.GroupNorm(epsilon=1e-4, name="gn_proj")(
              residual
          )
  
      if self.bottleneck:
          x = bit_resnet.StdConv(
              features, (1, 1), use_bias=False, name="conv1"
          )(x)
          x = nn.GroupNorm(epsilon=1e-4, name="gn1")(x)
          x = nn.relu(x)
  
      # Axial block that is replacing the 3x3 Convs in the ResNet residual unit.
      # Row  attention:
      x = AxialSelfAttention(
          attention_axis=1,
          axial_attention_configs=self.axial_attention_configs,
      )(x)
      # Column attention:
      x = AxialSelfAttention(
          attention_axis=2,
          axial_attention_configs=self.axial_attention_configs,
      )(x)
      if self.strides == (2, 2):
          x = nn.avg_pool(x, (2, 2), strides=(2, 2), padding="SAME")
      x = nn.relu(x)
  
      last_kernel = (1, 1) if self.bottleneck else (3, 3)
      x = bit_resnet.StdConv(
          nout, last_kernel, use_bias=False, name="conv3"
      )(x)
      x = nn.GroupNorm(
          epsilon=1e-4, name="gn3", scale_init=nn.initializers.zeros
      )(x)
      x = nn.relu(residual + x)
  
      return x
  
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 AxialResidualUnit obj's underlying member method code with example usage: obj.__call__(x) and returns: x