class_docstr: ResNet Stage: one or more stacked ResNet blocks.
  
  Attributes:
    block_size: Number of ResNet blocks to stack.
    nout: Number of features.
    first_stride: Downsampling stride.
    first_dilation: Kernel dilation.
    bottleneck: If True, the bottleneck block is used.
    gn_num_groups: Number of groups in group norm layer.
class_docstr_tok: ["ResNet","Stage",":","one","or","more","stacked","ResNet","blocks",".","Attributes",":","block_size",":","Number","of","ResNet","blocks","to","stack",".","nout",":","Number","of","features",".","first_stride",":","Downsampling","stride",".","first_dilation",":","Kernel","dilation",".","bottleneck",":","If","True",",","the","bottleneck","block","is","used",".","gn_num_groups",":","Number","of","groups","in","group","norm","layer","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      x = ResidualUnit(
          self.nout,
          strides=self.first_stride,
          dilation=self.first_dilation,
          bottleneck=self.bottleneck,
          gn_num_groups=self.gn_num_groups,
          name="unit1",
      )(x)
      for i in range(1, self.block_size):
          x = ResidualUnit(
              self.nout,
              strides=(1, 1),
              bottleneck=self.bottleneck,
              gn_num_groups=self.gn_num_groups,
              name=f"unit{i + 1}",
          )(x)
      return x
  
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 ResNetStage obj's underlying member method code with example usage: obj.__call__(x) and returns: x