class_docstr: Axial ResNet.
  
  Attributes:
    num_outputs: Num output classes. If None, a dict of intermediate feature
      maps is returned.
    width_factor: Width multiplier for each of the ResNet stages.
    num_layers: Number of layers (see `BLOCK_SIZE_OPTIONS` for stage
      configurations).
class_docstr_tok: ["Axial","ResNet",".","Attributes",":","num_outputs",":","Num","output","classes",".","If","None",",","a","dict","of","intermediate","feature","maps","is","returned",".","width_factor",":","Width","multiplier","for","each","of","the","ResNet","stages",".","num_layers",":","Number","of","layers","(","see","`","BLOCK_SIZE_OPTIONS","`","for","stage","configurations",")","."]
docstr: Applies the AxialResNet model to the inputs.
  
  Args:
    x: Inputs to the model.
    train: Unused.
    debug: Unused.
  
  Returns:
     Un-normalized logits if `num_outputs` is provided, a dictionary with
     representations otherwise.
docstr_tok: ["Applies","the","AxialResNet","model","to","the","inputs",".","Args",":","x",":","Inputs","to","the","model",".","train",":","Unused",".","debug",":","Unused",".","Returns",":","Un-normalized","logits","if","`","num_outputs","`","is","provided",",","a","dictionary","with","representations","otherwise","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool = True, debug: bool = False
  ) -> Union[jnp.ndarray, Dict[str, jnp.ndarray]]:
      """Applies the AxialResNet model to the inputs.
  
      Args:
        x: Inputs to the model.
        train: Unused.
        debug: Unused.
  
      Returns:
         Un-normalized logits if `num_outputs` is provided, a dictionary with
         representations otherwise.
      """
      del train
      del debug
      blocks, bottleneck = bit_resnet.BLOCK_SIZE_OPTIONS[
          self.num_layers
      ]
      width = int(64 * self.width_factor)
  
      # Root block.
      x = bit_resnet.StdConv(
          width, (7, 7), (2, 2), use_bias=False, name="conv_root"
      )(x)
      x = nn.GroupNorm(epsilon=1e-4, name="gn_root")(x)
      x = nn.relu(x)
      x = nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")
  
      # Stages.
      x = AxialResNetStage(
          blocks[0],
          width,
          axial_attention_configs=self.axial_attention_configs,
          first_stride=(1, 1),
          bottleneck=bottleneck,
          name="block1",
      )(x)
      for i, block_size in enumerate(blocks[1:], 1):
          x = AxialResNetStage(
              block_size,
              width * 2**i,
              axial_attention_configs=self.axial_attention_configs,
              first_stride=(2, 2),
              bottleneck=bottleneck,
              name=f"block{i + 1}",
          )(x)
  
      # Head.
      x = jnp.mean(x, axis=(1, 2))
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      return nn.Dense(
          self.num_outputs,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(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 AxialResNet obj's underlying member method code with example usage: obj.__call__(x) and returns: nn