.    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
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 J