class_docstr: A ResNet class that is similar to torchvision's with changes.
  
  - There are now 3 "stem" convolutions as opposed to 1, with an average pool
  instead of a max pool.
  - Performs anti-aliasing strided convolutions, where an avgpool is
  prepended to convolutions with stride > 1 - The final pooling layer is a
  QKV attention instead of an average pool.
  
  Attributes:
    features: Number of features.
    out_features: Number of output features. If None, return resnet feature-map.
    num_layers: Number of layers for each block.
    num_heads: Number of heads.
class_docstr_tok: ["A","ResNet","class","that","is","similar","to","torchvision","'s","with","changes",".","-","There","are","now","3","``","stem","''","convolutions","as","opposed","to","1",",","with","an","average","pool","instead","of","a","max","pool",".","-","Performs","anti-aliasing","strided","convolutions",",","where","an","avgpool","is","prepended","to","convolutions","with","stride",">","1","-","The","final","pooling","layer","is","a","QKV","attention","instead","of","an","average","pool",".","Attributes",":","features",":","Number","of","features",".","out_features",":","Number","of","output","features",".","If","None",",","return","resnet","feature-map",".","num_layers",":","Number","of","layers","for","each","block",".","num_heads",":","Number","of","heads","."]
code: def setup(self):
      # The 3-layer stem.
      self.conv1 = nn.Conv(
          self.features // 2,
          kernel_size=(3, 3),
          strides=(2, 2),
          padding=[(1, 1), (1, 1)],
          use_bias=False,
          name="conv1",
      )
      self.bn1 = nn.BatchNorm(use_running_average=True, name="bn1")
      self.conv2 = nn.Conv(
          self.features // 2,
          kernel_size=(3, 3),
          padding=[(1, 1), (1, 1)],
          use_bias=False,
          name="conv2",
      )
      self.bn2 = nn.BatchNorm(use_running_average=True, name="bn2")
      self.conv3 = nn.Conv(
          self.features,
          kernel_size=(3, 3),
          padding=[(1, 1), (1, 1)],
          use_bias=False,
          name="conv3",
      )
      self.bn3 = nn.BatchNorm(use_running_average=True, name="bn3")
  
      # Residual layers.
      self.layer1 = ResNetStage(
          self.features, self.num_layers[0], name="layer1"
      )
      self.layer2 = ResNetStage(
          self.features * 2, self.num_layers[1], stride=2, name="layer2"
      )
      self.layer3 = ResNetStage(
          self.features * 4, self.num_layers[2], stride=2, name="layer3"
      )
      self.layer4 = ResNetStage(
          self.features * 8, self.num_layers[3], stride=2, name="layer4"
      )
      if self.out_features is not None:
          self.attnpool = AttentionPool(
              self.num_heads, self.out_features, name="attnpool"
          )
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the ModifiedResNet obj's underlying member method code with example usage: obj.setup() without return types