'Create a diagonal gaussian distribution.
Args:
batch_size: The size of the batch, i.e. 0th dim in 2D tensor of samples.
z_size: The dimension of the distribution, i.e. 1st dim in 2D tensor.
mean: The N-D mean of the distribution.
logvar: The N-D log variance of the diagonal distribution.'

def __init__(self, batch_size, z_size, mean, logvar):

    size__xz = [None, z_size]
    self.mean = mean
    self.logvar = logvar
    self.noise = noise = tf.random_normal(tf.shape(logvar))
    self.sample = (mean + (tf.exp((0.5 * logvar)) * noise))
    mean.set_shape(size__xz)
    logvar.set_shape(size__xz)
    self.sample.set_shape(size__xz)
