human
I am working on the file "Sessions/Session14/Day3/IntroductionToVariationalAutoencoders.ipynb".
The first task is:
STEP 3: Modify the AE class into a VAE class!

A VAE has a few more input parameters than an AE. We need a sample size, which determines how many samples we draw from $q_\phi(z|x)$ for evaluating the ELBO. We also need a $\sigma_\epsilon$ to characterize the likelihood, $p_\psi(x|z)=\mathcal{G}(x',\sigma_\epsilon)$.
The prior is fixed. We can define it in the beginning, when we initialize the VAE.
We need methods to compute the variational posterior, the likelihood, the KL-divergence and the ELBO.
Can you write Python code for it?
 ... Write me a Python code to compute the variational posterior, the likelihood, the KL-divergence and the ELBO?
assistant
#Here is Python code to compute the variational posterior, the likelihood, the KL-divergence and the ELBO.

class VariationalAutoencoder(nn.Module):
    #TASK: add parameters mentioned in point 1. 
    def __init__(self, ??, ??):
        super(VariationalAutoencoder, self).__init__()
        self.encoder = VAEEncoder()
        self.decoder = VAEDecoder()
        self."insert name of param1 here" = ??
        self."insert name of param2 here" = ??
        #TASK: Use the Normal class to define the prior (a standard normal distribution), p(z)
        self.prior       = ??

    def get_q(self,x):
        #TASK: write a method that computes q(z,x)
        #HINT: use the Normal class we imported above
        mu, std = ??
        self.q  = ??
        return True

    def sample_q(self):
        #TASK: write a method that samples from q
        #HINT: use rsample to apply the reparameterization trick
        #HINT: rsample takes a list or array of shapes as input, e.g. [sample_size]
        z_sample = ??
        return z_sample

        #TASK (Optional): write a method that allows to change how many samples are drawn from q(z|x)
    def change_sample_size(self,??):
        ?? = ??
        return True

    def get_avg_log_likelihood(self,recons,x):
        #TASK: Write a method that returns the first term in the ELBO (this method should define the likelihood and evaluate the average log likelihood of the reconstruction)
        #HINT: Pay attention to shapes. The function should return an average log likelihood (a single number) for every data point in the batch.
        #HINT: The output shape of Normal(mu, sigma).log_prob() is a little unintuitive. If mu or sigma are N-dimensional, it returns N results (applies N independent Gaussians). 
        #HINT: You need to average over samples from q to obtain the final result.
        ll    = ??
        log_p = ??
        log_p = ??
        return ??

    def stochastic_kl_divergence(self,z_sample):
        #TASK: Write a method that computes the kl-divergence between q(z|x) and p(z) 
        #HINT: Pay attention to shape
        return ??

    def forward(self, x):
       ...