class_docstr: Flax Unet Mid-Block module.
  
  Parameters:
      in_channels (:obj:`int`):
          Input channels
      dropout (:obj:`float`, *optional*, defaults to 0.0):
          Dropout rate
      num_layers (:obj:`int`, *optional*, defaults to 1):
          Number of Resnet layer block
      resnet_groups (:obj:`int`, *optional*, defaults to `32`):
          The number of groups to use for the Resnet and Attention block group norm
      num_attention_heads (:obj:`int`, *optional*, defaults to `1`):
          Number of attention heads for each attention block
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
class_docstr_tok: ["Flax","Unet","Mid-Block","module",".","Parameters",":","in_channels","(",":","obj",":","`","int","`",")",":","Input","channels","dropout","(",":","obj",":","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","Dropout","rate","num_layers","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","1",")",":","Number","of","Resnet","layer","block","resnet_groups","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","32","`",")",":","The","number","of","groups","to","use","for","the","Resnet","and","Attention","block","group","norm","num_attention_heads","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","1","`",")",":","Number","of","attention","heads","for","each","attention","block","dtype","(",":","obj",":","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","jnp.float32",")",":","Parameters","`","dtype","`"]
code: def setup(self):
      resnet_groups = (
          self.resnet_groups
          if self.resnet_groups is not None
          else min(self.in_channels // 4, 32)
      )
  
      # there is always at least one resnet
      resnets = [
          FlaxResnetBlock2D(
              in_channels=self.in_channels,
              out_channels=self.in_channels,
              dropout=self.dropout,
              groups=resnet_groups,
              dtype=self.dtype,
          )
      ]
  
      attentions = []
  
      for _ in range(self.num_layers):
          attn_block = FlaxAttentionBlock(
              channels=self.in_channels,
              num_head_channels=self.num_attention_heads,
              num_groups=resnet_groups,
              dtype=self.dtype,
          )
          attentions.append(attn_block)
  
          res_block = FlaxResnetBlock2D(
              in_channels=self.in_channels,
              out_channels=self.in_channels,
              dropout=self.dropout,
              groups=resnet_groups,
              dtype=self.dtype,
          )
          resnets.append(res_block)
  
      self.resnets = resnets
      self.attentions = attentions
  
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 FlaxUNetMidBlock2D obj's underlying member method code with example usage: obj.setup() without return types