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 paper.   dtype (jnp.dtype, optional, defaults to jnp.float32) —
The dtype of the parameters.    Flax implementation of a VAE model with KL loss for decoding latent representations. This model inherits from FlaxModelMixi