.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/cvae.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_cvae.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_cvae.py:


Example: Conditional Variational Autoencoder in Flax
====================================================

This example trains a *Conditional Variational Autoencoder* (CVAE) [1] on the MNIST data
using Flax' neural network API. The implementation can be found here:
https://github.com/pyro-ppl/numpyro/tree/master/examples/cvae-flax

The model is a port of Pyro's excellent CVAE example which describes the model as well as the data in detail:
https://pyro.ai/examples/cvae.html

The model first trains a baseline to predict an entire MNIST image from a single quadrant of it
(i.e., input is one quadrant of an image, output is the entire image (not the other three quadrants)).
Then, in a second model, the generation/prior/recognition nets of the CVAE are trained while keeping the model
parameters of the baseline fixed/frozen. We use Optax' `multi_transform` to apply different gradient transformations
to the trainable parameters and the frozen parameters.


.. image:: ../_static/img/examples/cvae.png
    :align: center

**References:**

    1. Kihyuk Sohn, Xinchen Yan, Honglak Lee (2015), "Learning Structured Output Representation using Deep
       Conditional Generative Models
       (https://papers.nips.cc/paper/5775-learning-structured-output-representation-using-deep-conditional-generative-models)


.. _sphx_glr_download_examples_cvae.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: cvae.ipynb <cvae.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: cvae.py <cvae.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: cvae.zip <cvae.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_