class_docstr: Bit ResNetV1.
  
  Attributes:
    num_outputs: Num output classes. If None, a dict of intermediate feature
      maps is returned
    gn_num_groups: Number groups in the group norm layer..
    width_factor: Width multiplier for each of the ResNet stages.
    num_layers: Number of layers (see `BLOCK_SIZE_OPTIONS` for stage
      configurations).
    max_output_stride: Defines the maximum output stride of the resnet.
      Typically, resnets output feature maps have sride 32. We can, however,
      lower that number by swapping strides with dilation in later stages. This
      is common in cases where stride 32 is too large, e.g., in dense prediciton
      tasks.
class_docstr_tok: ["Bit","ResNetV1",".","Attributes",":","num_outputs",":","Num","output","classes",".","If","None",",","a","dict","of","intermediate","feature","maps","is","returned","gn_num_groups",":","Number","groups","in","the","group","norm","layer","..","width_factor",":","Width","multiplier","for","each","of","the","ResNet","stages",".","num_layers",":","Number","of","layers","(","see","`","BLOCK_SIZE_OPTIONS","`","for","stage","configurations",")",".","max_output_stride",":","Defines","the","maximum","output","stride","of","the","resnet",".","Typically",",","resnets","output","feature","maps","have","sride","32",".","We","can",",","however",",","lower","that","number","by","swapping","strides","with","dilation","in","later","stages",".","This","is","common","in","cases","where","stride","32","is","too","large",",","e.g.",",","in","dense","prediciton","tasks","."]
docstr: Applies the Bit ResNet 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","Bit","ResNet","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 Bit ResNet 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
      if self.max_output_stride not in [4, 8, 16, 32]:
          raise ValueError(
              "Only supports output strides of [4, 8, 16, 32]"
          )
  
      blocks, bottleneck = BLOCK_SIZE_OPTIONS[self.num_layers]
  
      width = int(64 * self.width_factor)
  
      # Root block.
      x = StdConv(
          width, (7, 7), (2, 2), use_bias=False, name="conv_root"
      )(x)
      x = nn.GroupNorm(
          num_groups=self.gn_num_groups, epsilon=1e-4, name="gn_root"
      )(x)
      x = nn.relu(x)
      x = nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")
      representations = {"stem": x}
  
      # Stages.
      x = ResNetStage(
          blocks[0],
          width,
          first_stride=(1, 1),
          bottleneck=bottleneck,
          gn_num_groups=self.gn_num_groups,
          name="block1",
      )(x)
      stride = 4
      for i, block_size in enumerate(blocks[1:], 1):
          max_stride_reached = self.max_output_stride <= stride
          x = ResNetStage(
              block_size,
              width * 2**i,
              first_stride=(2, 2) if not max_stride_reached else (1, 1),
              first_dilation=(2, 2) if max_stride_reached else (1, 1),
              bottleneck=bottleneck,
              gn_num_groups=self.gn_num_groups,
              name=f"block{i + 1}",
          )(x)
          if not max_stride_reached:
              stride *= 2
          representations[f"stage_{i + 1}"] = x
  
      if self.num_outputs:
          # Head.
          x = jnp.mean(x, axis=(1, 2))
          x = nn_layers.IdentityLayer(name="pre_logits")(x)
          x = nn.Dense(
              self.num_outputs,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )(x)
          return x
      return representations
  
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the BitResNet obj's underlying member method code with example usage: obj.__call__(x, train, debug) and returns: representations, x