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 __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      def stem(x):
          for conv, bn in [
              (self.conv1, self.bn1),
              (self.conv2, self.bn2),
              (self.conv3, self.bn3),
          ]:
              x = nn.relu(bn(conv(x)))
          x = nn.avg_pool(x, (2, 2), (2, 2))
          return x
  
      x = stem(x)
      x = self.layer1(x)
      x = self.layer2(x)
      x = self.layer3(x)
      x = feature_map = self.layer4(x)
  
      if self.out_features is not None:
          x = self.attnpool(x)
  
      return x, feature_map
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the ModifiedResNet obj's underlying member method code with example usage: obj.__call__(x) and returns: x,  feature_map, x