class_docstr: Flax Implementation of VAE Decoder.
  
  This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)
  subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to
  general usage and behavior.
  
  Finally, this model supports inherent JAX features such as:
  - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
  - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
  - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
  - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
  
  Parameters:
      in_channels (:obj:`int`, *optional*, defaults to 3):
          Input channels
      out_channels (:obj:`int`, *optional*, defaults to 3):
          Output channels
      up_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):
          UpDecoder block type
      block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`):
          Tuple containing the number of output channels for each block
      layers_per_block (:obj:`int`, *optional*, defaults to `2`):
          Number of Resnet layer for each block
      norm_num_groups (:obj:`int`, *optional*, defaults to `32`):
          norm num group
      act_fn (:obj:`str`, *optional*, defaults to `silu`):
          Activation function
      double_z (:obj:`bool`, *optional*, defaults to `False`):
          Whether to double the last output channels
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          parameters `dtype`
class_docstr_tok: ["Flax","Implementation","of","VAE","Decoder",".","This","model","is","a","Flax","Linen","[","flax.linen.Module","]","(","https",":","\/\/flax.readthedocs.io\/en\/latest\/flax.linen.html","#","module",")","subclass",".","Use","it","as","a","regular","Flax","linen","Module","and","refer","to","the","Flax","documentation","for","all","matter","related","to","general","usage","and","behavior",".","Finally",",","this","model","supports","inherent","JAX","features","such","as",":","-","[","Just-In-Time","(","JIT",")","compilation","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","just-in-time-compilation-jit",")","-","[","Automatic","Differentiation","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","automatic-differentiation",")","-","[","Vectorization","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","vectorization-vmap",")","-","[","Parallelization","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","parallelization-pmap",")","Parameters",":","in_channels","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","3",")",":","Input","channels","out_channels","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","3",")",":","Output","channels","up_block_types","(",":","obj",":","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","UpDecoderBlock2D",")","`",")",":","UpDecoder","block","type","block_out_channels","(",":","obj",":","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","64",",",")","`",")",":","Tuple","containing","the","number","of","output","channels","for","each","block","layers_per_block","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","2","`",")",":","Number","of","Resnet","layer","for","each","block","norm_num_groups","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","32","`",")",":","norm","num","group","act_fn","(",":","obj",":","`","str","`",",","*","optional","*",",","defaults","to","`","silu","`",")",":","Activation","function","double_z","(",":","obj",":","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","to","double","the","last","output","channels","dtype","(",":","obj",":","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","jnp.float32",")",":","parameters","`","dtype","`"]
code: def setup(self):
      block_out_channels = self.block_out_channels
  
      # z to block_in
      self.conv_in = nn.Conv(
          block_out_channels[-1],
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          dtype=self.dtype,
      )
  
      # middle
      self.mid_block = FlaxUNetMidBlock2D(
          in_channels=block_out_channels[-1],
          resnet_groups=self.norm_num_groups,
          num_attention_heads=None,
          dtype=self.dtype,
      )
  
      # upsampling
      reversed_block_out_channels = list(reversed(block_out_channels))
      output_channel = reversed_block_out_channels[0]
      up_blocks = []
      for i, _ in enumerate(self.up_block_types):
          prev_output_channel = output_channel
          output_channel = reversed_block_out_channels[i]
  
          is_final_block = i == len(block_out_channels) - 1
  
          up_block = FlaxUpDecoderBlock2D(
              in_channels=prev_output_channel,
              out_channels=output_channel,
              num_layers=self.layers_per_block + 1,
              resnet_groups=self.norm_num_groups,
              add_upsample=not is_final_block,
              dtype=self.dtype,
          )
          up_blocks.append(up_block)
          prev_output_channel = output_channel
  
      self.up_blocks = up_blocks
  
      # end
      self.conv_norm_out = nn.GroupNorm(
          num_groups=self.norm_num_groups, epsilon=1e-6
      )
      self.conv_out = nn.Conv(
          self.out_channels,
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          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 FlaxDecoder obj's underlying member method code with example usage: obj.setup() without return types