class_docstr: Flax 2D downsizing block
  
  Parameters:
      in_channels (:obj:`int`):
          Input channels
      out_channels (:obj:`int`):
          Output channels
      dropout (:obj:`float`, *optional*, defaults to 0.0):
          Dropout rate
      num_layers (:obj:`int`, *optional*, defaults to 1):
          Number of attention blocks layers
      add_downsample (:obj:`bool`, *optional*, defaults to `True`):
          Whether to add downsampling layer before each final output
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
class_docstr_tok: ["Flax","2D","downsizing","block","Parameters",":","in_channels","(",":","obj",":","`","int","`",")",":","Input","channels","out_channels","(",":","obj",":","`","int","`",")",":","Output","channels","dropout","(",":","obj",":","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","Dropout","rate","num_layers","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","1",")",":","Number","of","attention","blocks","layers","add_downsample","(",":","obj",":","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","to","add","downsampling","layer","before","each","final","output","dtype","(",":","obj",":","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","jnp.float32",")",":","Parameters","`","dtype","`"]
code: def setup(self):
      resnets = []
  
      for i in range(self.num_layers):
          in_channels = (
              self.in_channels if i == 0 else self.out_channels
          )
  
          res_block = FlaxResnetBlock2D(
              in_channels=in_channels,
              out_channels=self.out_channels,
              dropout_prob=self.dropout,
              dtype=self.dtype,
          )
          resnets.append(res_block)
      self.resnets = resnets
  
      if self.add_downsample:
          self.downsamplers_0 = FlaxDownsample2D(
              self.out_channels, dtype=self.dtype
          )
  
imports: import flax.linen
  import jax.numpy
  from .attention_flax import FlaxTransformer2DModel
  from .resnet_flax import FlaxDownsample2D, FlaxResnetBlock2D, FlaxUpsample2D
label_desc: Use image node_id 1 for calling the FlaxDownBlock2D obj's underlying member method code with example usage: obj.setup() without return types