text: import argparse
  
  import torch
  import torch.nn as nn
  from torch.nn import functional as F
  import matplotlib.pyplot as plt
  from torchvision.utils import make_grid
  
  from datasets.bmnist import bmnist
  
  from
  
  
  class Encoder(nn.Module):
  
      def __init__(self, hidden_dim=500, z_dim=20):
          super().__init__()
  
          # FIXME: Get input dimension
          self.mean_en1  = nn.Linear(784, hidden_dim)
          self.mean_en2  = nn.Linear(hidden_dim, z_dim)
  
          self.std_en1   = nn.Linear(784, hidden_dim)
          self.std_en2   = nn.Linear(hidden_dim, z_dim)
  
      def forward(self, input):
          """
          Perform forward pass of encoder.
  
          Returns mean and std with shape [batch_size, z_dim]. Make sure
          that any constraints are enforced.
          """
  
          mean = self.mean_en2(F.relu(self.mean_en1(input)))
          std  = self.std_en2(F.relu(self.std_en1(input)))
  
          return mean, std
  
  
  class Decoder(nn.Module):
  
      def __init__(self, hidden_dim=500, z_dim=20):
          super().__init__()
  
          self.dn1 = nn.Linear(z_dim, hidden_dim)
          self.dn2 = nn.Linaer(hidden_dim, 784)
  
      def forward(self, input):
          """
          Perform forward pass of encoder.
  
          Returns mean with shape [batch_size, 784].
          """
  
          mean = self.dn2(F.relu(self.dn1(input)))
          mean = torch.sigmoid(mean)
  
          return mean
  
  
  class VAE(nn.Module):
  
      def __init__(self, hidden_dim=500, z_dim=20, device='cpu'):
          super().__init__()
  
          self.z_dim = z_dim
          self.encoder = Encoder(hidden_dim, z_dim)
          self.decoder = Decoder(hidden_dim, z_dim)
  
          self.device  = device
          self.to(device)
  
      def forward(self, input):
          """
          Given input, perform an encoding and decoding step and return the
          negative average elbo for the given batch.
          """
          input = input.view(-1, 784).to(self.device)
  
          mean, std             = self.encoder(input)
          z                     = self.reparameterize(mean, std)
          reconstructed         = self.decoder(z)
          average_negative_elbo = loss_function(input, reconstructed, mean, std)
  
          return average_negative_elbo
  
  
      def loss_function(input, reconstructed, mu, logvar):
          """
          ELBO: Cross Entropy + KL
          """
          cross_entropy_error = F.binary_cross_entropy(reconstructed, input, reduction='sum')
          KL_term             = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
  
          return cross_entropy_error + KL_term
  
      def reparameterize(self, mu, logvar):
          """
          Reparameterization trick
          """
          std = torch.exp(0.5 * logvar)
          eps = torch.randn_like(std)
  
          return eps.mul(std).add_(mu)
  
          average_negative_elbo = None
          raise NotImplementedError()
          return average_negative_elbo
  
      def sample(self, n_samples):
          """
          Sample n_samples from the model. Return both the sampled images
          (from bernoulli) and the means for these bernoullis (as these are
          used to plot the data manifold).
          """
          sampled_ims, im_means = None, None
          raise NotImplementedError()
  
          return sampled_ims, im_means
  
  
  def epoch_iter(model, data, optimizer):
      """
      Perform a single epoch for either the training or validation.
      use model.training to determine if in 'training mode' or not.
  
      Returns the average elbo for the complete epoch.
      """
      average_epoch_elbo = None
      raise NotImplementedError()
  
      return average_epoch_elbo
  
  
  def run_epoch(model, data, optimizer):
      """
      Run a train and validation epoch and return average elbo for each.
      """
      traindata, valdata = data
  
      model.train()
      train_elbo = epoch_iter(model, traindata, optimizer)
  
      model.eval()
      val_elbo = epoch_iter(model, valdata, optimizer)
  
      return train_elbo, val_elbo
  
  
  def save_elbo_plot(train_curve, val_curve, filename):
      plt.figure(figsize=(12, 6))
      plt.plot(train_curve, label='train elbo')
      plt.plot(val_curve, label='validation elbo')
      plt.legend()
      plt.xlabel('epochs')
      plt.ylabel('ELBO')
      plt.tight_layout()
      plt.savefig(filename)
  
  
  def main():
      data = bmnist()[:2]  # ignore test split
      model = VAE(z_dim=ARGS.zdim)
      optimizer = torch.optim.Adam(model.parameters())
  
      train_curve, val_curve = [], []
      for epoch in range(ARGS.epochs):
          elbos = run_epoch(model, data, optimizer)
          train_elbo, val_elbo = elbos
          train_curve.append(train_elbo)
          val_curve.append(val_elbo)
          print(f"[Epoch {epoch}] train elbo: {train_elbo} val_elbo: {val_elbo}")
  
          # --------------------------------------------------------------------
          #  Add functionality to plot samples from model during training.
          #  You can use the make_grid functioanlity that is already imported.
          # --------------------------------------------------------------------
  
      # --------------------------------------------------------------------
      #  Add functionality to plot plot the learned data manifold after
      #  if required (i.e., if zdim == 2). You can use the make_grid
      #  functionality that is already imported.
      # --------------------------------------------------------------------
  
      save_elbo_plot(train_curve, val_curve, 'elbo.pdf')
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument('--epochs', default=40, type=int,
                          help='max number of epochs')
      parser.add_argument('--zdim', default=20, type=int,
                          help='dimensionality of latent space')
  
      ARGS = parser.parse_args()
  
      main()
  