class_docstr: Flax implementation of 2D Resnet Block.
  
  Args:
      in_channels (`int`):
          Input channels
      out_channels (`int`):
          Output channels
      dropout (:obj:`float`, *optional*, defaults to 0.0):
          Dropout rate
      groups (:obj:`int`, *optional*, defaults to `32`):
          The number of groups to use for group norm.
      use_nin_shortcut (:obj:`bool`, *optional*, defaults to `None`):
          Whether to use `nin_shortcut`. This activates a new layer inside ResNet block
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
class_docstr_tok: ["Flax","implementation","of","2D","Resnet","Block",".","Args",":","in_channels","(","`","int","`",")",":","Input","channels","out_channels","(","`","int","`",")",":","Output","channels","dropout","(",":","obj",":","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","Dropout","rate","groups","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","32","`",")",":","The","number","of","groups","to","use","for","group","norm",".","use_nin_shortcut","(",":","obj",":","`","bool","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Whether","to","use","`","nin_shortcut","`",".","This","activates","a","new","layer","inside","ResNet","block","dtype","(",":","obj",":","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","jnp.float32",")",":","Parameters","`","dtype","`"]
code: def setup(self):
      out_channels = (
          self.in_channels
          if self.out_channels is None
          else self.out_channels
      )
  
      self.norm1 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6)
      self.conv1 = nn.Conv(
          out_channels,
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          dtype=self.dtype,
      )
  
      self.norm2 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6)
      self.dropout_layer = nn.Dropout(self.dropout)
      self.conv2 = nn.Conv(
          out_channels,
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          dtype=self.dtype,
      )
  
      use_nin_shortcut = (
          self.in_channels != out_channels
          if self.use_nin_shortcut is None
          else self.use_nin_shortcut
      )
  
      self.conv_shortcut = None
      if use_nin_shortcut:
          self.conv_shortcut = nn.Conv(
              out_channels,
              kernel_size=(1, 1),
              strides=(1, 1),
              padding="VALID",
              dtype=self.dtype,
          )
  
imports: import math
  from functools import partial
  from typing import Tuple
  import flax
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict
  from ..configuration_utils import ConfigMixin, flax_register_to_config
  from ..utils import BaseOutput
  from .modeling_flax_utils import FlaxModelMixin
label_desc: Use image node_id 1 for calling the FlaxResnetBlock2D obj's underlying member method code with example usage: obj.setup() without return types