class_docstr: Flax implementation of a VAE model with KL loss for decoding latent representations.
  
  This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods
  implemented for all models (such as downloading or saving).
  
  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 its
  general usage and behavior.
  
  Inherent JAX features such as the following are supported:
  
  - [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 (`int`, *optional*, defaults to 3):
          Number of channels in the input image.
      out_channels (`int`, *optional*, defaults to 3):
          Number of channels in the output.
      down_block_types (`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):
          Tuple of downsample block types.
      up_block_types (`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):
          Tuple of upsample block types.
      block_out_channels (`Tuple[str]`, *optional*, defaults to `(64,)`):
          Tuple of block output channels.
      layers_per_block (`int`, *optional*, defaults to `2`):
          Number of ResNet layer for each block.
      act_fn (`str`, *optional*, defaults to `silu`):
          The activation function to use.
      latent_channels (`int`, *optional*, defaults to `4`):
          Number of channels in the latent space.
      norm_num_groups (`int`, *optional*, defaults to `32`):
          The number of groups for normalization.
      sample_size (`int`, *optional*, defaults to 32):
          Sample input size.
      scaling_factor (`float`, *optional*, defaults to 0.18215):
          The component-wise standard deviation of the trained latent space computed using the first batch of the
          training set. This is used to scale the latent space to have unit variance when training the diffusion
          model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
          diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
          / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
          Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
      dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
          The `dtype` of the parameters.
class_docstr_tok: ["Flax","implementation","of","a","VAE","model","with","KL","loss","for","decoding","latent","representations",".","This","model","inherits","from","[","`","FlaxModelMixin","`","]",".","Check","the","superclass","documentation","for","it","'s","generic","methods","implemented","for","all","models","(","such","as","downloading","or","saving",")",".","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","its","general","usage","and","behavior",".","Inherent","JAX","features","such","as","the","following","are","supported",":","-","[","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","(","`","int","`",",","*","optional","*",",","defaults","to","3",")",":","Number","of","channels","in","the","input","image",".","out_channels","(","`","int","`",",","*","optional","*",",","defaults","to","3",")",":","Number","of","channels","in","the","output",".","down_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","DownEncoderBlock2D",")","`",")",":","Tuple","of","downsample","block","types",".","up_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","UpDecoderBlock2D",")","`",")",":","Tuple","of","upsample","block","types",".","block_out_channels","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","64",",",")","`",")",":","Tuple","of","block","output","channels",".","layers_per_block","(","`","int","`",",","*","optional","*",",","defaults","to","`","2","`",")",":","Number","of","ResNet","layer","for","each","block",".","act_fn","(","`","str","`",",","*","optional","*",",","defaults","to","`","silu","`",")",":","The","activation","function","to","use",".","latent_channels","(","`","int","`",",","*","optional","*",",","defaults","to","`","4","`",")",":","Number","of","channels","in","the","latent","space",".","norm_num_groups","(","`","int","`",",","*","optional","*",",","defaults","to","`","32","`",")",":","The","number","of","groups","for","normalization",".","sample_size","(","`","int","`",",","*","optional","*",",","defaults","to","32",")",":","Sample","input","size",".","scaling_factor","(","`","float","`",",","*","optional","*",",","defaults","to","0.18215",")",":","The","component-wise","standard","deviation","of","the","trained","latent","space","computed","using","the","first","batch","of","the","training","set",".","This","is","used","to","scale","the","latent","space","to","have","unit","variance","when","training","the","diffusion","model",".","The","latents","are","scaled","with","the","formula","`","z","=","z","*","scaling_factor","`","before","being","passed","to","the","diffusion","model",".","When","decoding",",","the","latents","are","scaled","back","to","the","original","scale","with","the","formula",":","`","z","=","1","\/","scaling_factor","*","z","`",".","For","more","details",",","refer","to","sections","4.3.2","and","D.1","of","the","[","High-Resolution","Image","Synthesis","with","Latent","Diffusion","Models","]","(","https",":","\/\/arxiv.org\/abs\/2112.10752",")","paper",".","dtype","(","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","`","jnp.float32","`",")",":","The","`","dtype","`","of","the","parameters","."]
code: def init_weights(self, rng: jax.Array) -> FrozenDict:
      # init input tensors
      sample_shape = (
          1,
          self.in_channels,
          self.sample_size,
          self.sample_size,
      )
      sample = jnp.zeros(sample_shape, dtype=jnp.float32)
  
      params_rng, dropout_rng, gaussian_rng = jax.random.split(rng, 3)
      rngs = {
          "params": params_rng,
          "dropout": dropout_rng,
          "gaussian": gaussian_rng,
      }
  
      return self.init(rngs, sample)["params"]
  
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 2 for calling the FlaxAutoencoderKL obj's underlying member method code with example usage: obj.init_weights(rng) and returns: self