class_docstr: Attention pooling layer.
  
  Attributes:
    features: Number of features.
    num_layers: Number of bottleneck blocks.
    stride: Stride in the Bottleneck module.
class_docstr_tok: ["Attention","pooling","layer",".","Attributes",":","features",":","Number","of","features",".","num_layers",":","Number","of","bottleneck","blocks",".","stride",":","Stride","in","the","Bottleneck","module","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      x = Bottleneck(self.features, self.stride, name="0")(x)
      for i in range(1, self.num_layers):
          x = Bottleneck(self.features, name=str(i))(x)
      return x
  
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 ResNetStage obj's underlying member method code with example usage: obj.__call__(x) and returns: x