class_docstr: ResNet architecture.
  
  Attributes:
    num_outputs: Num output classes. If None, a dict of intermediate feature
      maps is returned.
    num_filters: Num filters.
    num_layers: Num layers.
    kernel_init: Kernel initialization.
    bias_init: Bias initialization.
    dtype: Data type, e.g. jnp.float32.
class_docstr_tok: ["ResNet","architecture",".","Attributes",":","num_outputs",":","Num","output","classes",".","If","None",",","a","dict","of","intermediate","feature","maps","is","returned",".","num_filters",":","Num","filters",".","num_layers",":","Num","layers",".","kernel_init",":","Kernel","initialization",".","bias_init",":","Bias","initialization",".","dtype",":","Data","type",",","e.g",".","jnp.float32","."]
docstr: Applies ResNet model to the inputs.
  
  Args:
    x: Inputs to the model.
    train: Whether it is training or not.
    debug: Whether the debug mode is enabled. debug=True enables model
      specific logging/storing some values using jax.host_callback.
  
  Returns:
     Un-normalized logits.
docstr_tok: ["Applies","ResNet","model","to","the","inputs",".","Args",":","x",":","Inputs","to","the","model",".","train",":","Whether","it","is","training","or","not",".","debug",":","Whether","the","debug","mode","is","enabled",".","debug=True","enables","model","specific","logging\/storing","some","values","using","jax.host_callback",".","Returns",":","Un-normalized","logits","."]
code: def __call__(
      self, x: jnp.ndarray, train: bool = False, debug: bool = False
  ) -> Union[jnp.ndarray, Dict[str, jnp.ndarray]]:
      """Applies ResNet model to the inputs.
  
      Args:
        x: Inputs to the model.
        train: Whether it is training or not.
        debug: Whether the debug mode is enabled. debug=True enables model
          specific logging/storing some values using jax.host_callback.
  
      Returns:
         Un-normalized logits.
      """
      if self.num_layers not in BLOCK_SIZE_OPTIONS:
          raise ValueError("Please provide a valid number of layers")
      block_sizes, bottleneck = BLOCK_SIZE_OPTIONS[self.num_layers]
      x = nn.Conv(
          self.num_filters,
          kernel_size=(7, 7),
          strides=(2, 2),
          padding=[(3, 3), (3, 3)],
          use_bias=False,
          dtype=self.dtype,
          name="stem_conv",
      )(x)
      x = nn.BatchNorm(
          use_running_average=not train,
          momentum=0.9,
          epsilon=1e-5,
          dtype=self.dtype,
          name="init_bn",
      )(x)
      x = nn_layers.IdentityLayer(name="init_relu")(nn.relu(x))
      x = nn.max_pool(
          x, (3, 3), strides=(2, 2), padding=[(1, 1), (1, 1)]
      )
      x = nn_layers.IdentityLayer(name="stem_pool")(x)
  
      residual_block = functools.partial(
          ResidualBlock, dtype=self.dtype, bottleneck=bottleneck
      )
      representations = {"stem": x}
      for i, block_size in enumerate(block_sizes):
          for j in range(block_size):
              strides = (2, 2) if i > 0 and j == 0 else (1, 1)
              filters = self.num_filters * 2**i
              x = residual_block(filters=filters, strides=strides)(
                  x, train
              )
          representations[f"stage_{i + 1}"] = x
  
      # Head.
      if self.num_outputs:
          x = jnp.mean(x, axis=(1, 2))
          x = nn_layers.IdentityLayer(name="pre_logits")(x)
          x = nn.Dense(
              self.num_outputs,
              kernel_init=self.kernel_init,
              bias_init=self.bias_init,
              dtype=self.dtype,
              name="output_projection",
          )(x)
          return x
      else:
          return representations
  
imports: import functools
  from typing import Tuple, Callable, Any, Optional, Union, Dict
  from absl import logging
  import flax
  import flax.linen
  from jax.nn import initializers
  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 ResNet obj's underlying member method code with example usage: obj.__call__(x, train, debug) and returns: x, representations