class_docstr: This class encapsulates the parameters (neural networks) and models & guides needed to train a
  semi-supervised variational auto-encoder on the MNIST image dataset
  
  :param output_size: size of the tensor representing the class label (10 for MNIST since
                      we represent the class labels as a one-hot vector with 10 components)
  :param input_size: size of the tensor representing the image (28*28 = 784 for our MNIST dataset
                     since we flatten the images and scale the pixels to be in [0,1])
  :param z_dim: size of the tensor representing the latent random variable z
                (handwriting style for our MNIST dataset)
  :param hidden_layers: a tuple (or list) of MLP layers to be used in the neural networks
                        representing the parameters of the distributions in our model
  :param use_cuda: use GPUs for faster training
  :param aux_loss_multiplier: the multiplier to use with the auxiliary loss
docstr: The model corresponds to the following generative process:
  p(z) = normal(0,I)              # handwriting style (latent)
  p(y|x) = categorical(I/10.)     # which digit (semi-supervised)
  p(x|y,z) = bernoulli(loc(y,z))   # an image
  loc is given by a neural network  `decoder`
  
  :param xs: a batch of scaled vectors of pixels from an image
  :param ys: (optional) a batch of the class labels i.e.
             the digit corresponding to the image(s)
  :return: None
docstr_tok: ["The","model","corresponds","to","the","following","generative","process",":","p","(","z",")","=","normal","(","0",",","I",")","#","handwriting","style","(","latent",")","p","(","y|x",")","=","categorical","(","I\/10",".",")","#","which","digit","(","semi-supervised",")","p","(","x|y",",","z",")","=","bernoulli","(","loc","(","y",",","z",")",")","#","an","image","loc","is","given","by","a","neural","network","`","decoder","`",":","param","xs",":","a","batch","of","scaled","vectors","of","pixels","from","an","image",":","param","ys",":","(","optional",")","a","batch","of","the","class","labels","i.e",".","the","digit","corresponding","to","the","image","(","s",")",":","return",":","None"]
code: def model(self, xs, ys=None):
      """
      The model corresponds to the following generative process:
      p(z) = normal(0,I)              # handwriting style (latent)
      p(y|x) = categorical(I/10.)     # which digit (semi-supervised)
      p(x|y,z) = bernoulli(loc(y,z))   # an image
      loc is given by a neural network  `decoder`
  
      :param xs: a batch of scaled vectors of pixels from an image
      :param ys: (optional) a batch of the class labels i.e.
                 the digit corresponding to the image(s)
      :return: None
      """
      # register this pytorch module and all of its sub-modules with pyro
      pyro.module("ss_vae", self)
  
      batch_size = xs.size(0)
      options = dict(dtype=xs.dtype, device=xs.device)
      with pyro.plate("data"):
          # sample the handwriting style from the constant prior distribution
          prior_loc = torch.zeros(batch_size, self.z_dim, **options)
          prior_scale = torch.ones(batch_size, self.z_dim, **options)
          zs = pyro.sample(
              "z", dist.Normal(prior_loc, prior_scale).to_event(1)
          )
  
          # if the label y (which digit to write) is supervised, sample from the
          # constant prior, otherwise, observe the value (i.e. score it against the constant prior)
          alpha_prior = torch.ones(
              batch_size, self.output_size, **options
          ) / (1.0 * self.output_size)
          ys = pyro.sample(
              "y", dist.OneHotCategorical(alpha_prior), obs=ys
          )
  
          # Finally, score the image (x) using the handwriting style (z) and
          # the class label y (which digit to write) against the
          # parametrized distribution p(x|y,z) = bernoulli(decoder(y,z))
          # where `decoder` is a neural network. We disable validation
          # since the decoder output is a relaxed Bernoulli value.
          loc = self.decoder([zs, ys])
          pyro.sample(
              "x",
              dist.Bernoulli(loc, validate_args=False).to_event(1),
              obs=xs,
          )
          # return the loc so we can visualize it later
          return loc
  
imports: import argparse
  import torch
  import torch.nn
  from utils.custom_mlp import MLP, Exp
  from utils.mnist_cached import MNISTCached, mkdir_p, setup_data_loaders
  from utils.vae_plots import mnist_test_tsne_ssvae, plot_conditional_samples_ssvae
  from visdom import Visdom
  import pyro
  import pyro.distributions
  from pyro.contrib.examples.util import print_and_log
  from pyro.infer import SVI, JitTrace_ELBO, JitTraceEnum_ELBO, Trace_ELBO, TraceEnum_ELBO, config_enumerate
  from pyro.optim import Adam
label_desc: Use image node_id 3 for calling the SSVAE obj's underlying member method code with example usage: obj.model(xs, ys) and returns: loc