class_docstr: Bottleneck layer of ResNet.
  
  Attributes:
    features: Number of features.
    stride: Stride of the down-sampled output.
    expansion: Expansion of feature dimension.
class_docstr_tok: ["Bottleneck","layer","of","ResNet",".","Attributes",":","features",":","Number","of","features",".","stride",":","Stride","of","the","down-sampled","output",".","expansion",":","Expansion","of","feature","dimension","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      conv1 = nn.Conv(
          self.features, (1, 1), use_bias=False, name="conv1"
      )
      bn1 = nn.BatchNorm(use_running_average=True, name="bn1")
  
      conv2 = nn.Conv(
          self.features,
          (3, 3),
          padding=[(1, 1), (1, 1)],
          use_bias=False,
          name="conv2",
      )
      bn2 = nn.BatchNorm(use_running_average=True, name="bn2")
  
      conv3 = nn.Conv(
          self.features * self.expansion,
          (1, 1),
          use_bias=False,
          name="conv3",
      )
      bn3 = nn.BatchNorm(use_running_average=True, name="bn3")
  
      out = nn.relu(bn1(conv1(x)))
      out = nn.relu(bn2(conv2(out)))
      out = nn.avg_pool(
          out, (self.stride, self.stride), (self.stride, self.stride)
      )
      out = bn3(conv3(out))
  
      downsample = (
          self.stride > 1
          or x.shape[-1] != self.features * self.expansion
      )
      if downsample:
          x = Shortcut(
              features=self.features * self.expansion,
              stride=self.stride,
              name="downsample",
          )(x)
  
      out += x
      out = nn.relu(out)
      return out
  
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 Bottleneck obj's underlying member method code with example usage: obj.__call__(x) and returns: out