text: 
  """
  PixelVAE: A Latent Variable Model for Natural Images
  Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taiga, Francesco Visin, David Vazquez, Aaron Courville
  """
  
  import os, sys
  sys.path.append(os.getcwd())
  
  N_GPUS = 2
  
  import tflib as lib
  import tflib.train_loop_cifar_filter_3_dim_2
  import tflib.ops.kl_unit_gaussian
  import tflib.ops.kl_gaussian_gaussian
  import tflib.ops.conv2d
  import tflib.ops.linear
  import tflib.ops.batchnorm
  import tflib.ops.embedding
  
  import tflib.cifar
  import tflib.cifar_256
  
  import numpy as np
  import tensorflow as tf
  import imageio
  from imageio import imsave
  
  import time
  import functools
  
  DATASET = 'cifar10'#'lsun_64'
  SETTINGS =  '32px_cifar' #64px_big' # mnist_256, 32px_small, 32px_big, 64px_small, 64px_big
  
  OUT_DIR = DATASET + '_results' + '_filter_3_dim_2'
  
  if not os.path.isdir(OUT_DIR):
     os.makedirs(OUT_DIR)
     print "Created directory {}".format(OUT_DIR)
  
  if SETTINGS == 'mnist_256':
      # two_level uses Enc1/Dec1 for the bottom level, Enc2/Dec2 for the top level
      # one_level uses EncFull/DecFull for the bottom (and only) level
      MODE = 'one_level'
  
      # Whether to treat pixel inputs to the model as real-valued (as in the 
      # original PixelCNN) or discrete (gets better likelihoods).
      EMBED_INPUTS = True
  
      # Turn on/off the bottom-level PixelCNN in Dec1/DecFull
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 32
      DIM_1        = 16
      DIM_2        = 32
      DIM_3        = 32
      DIM_4        = 64
      LATENT_DIM_2 = 128
  
      ALPHA1_ITERS = 5000
      ALPHA2_ITERS = 5000
      KL_PENALTY = 1.0
      BETA_ITERS = 1000
  
      # In Dec2, we break each spatial location into N blocks (analogous to channels
      # in the original PixelCNN) and model each spatial location autoregressively
      # as P(x)=P(x0)*P(x1|x0)*P(x2|x0,x1)... In my experiments values of N > 1
      # actually hurt performance. Unsure why; might be a bug.
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 2*500,
          'stop_after': 500*500,
          'callback_every': 10*500
      }
  
      LR = 1e-3
  
      LR_DECAY_AFTER = TIMES['stop_after']
      LR_DECAY_FACTOR = 1.
  
      BATCH_SIZE = 100
      N_CHANNELS = 1
      HEIGHT = 28
      WIDTH = 28
  
      # These aren't actually used for one-level models but some parts
      # of the code still depend on them being defined.
      LATENT_DIM_1 = 64
      LATENTS1_HEIGHT = 7
      LATENTS1_WIDTH = 7
  
  elif SETTINGS == '32px_small':
      MODE = 'two_level'
  
      EMBED_INPUTS = True
  
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 128
      DIM_1        = 64
      DIM_2        = 128
      DIM_3        = 256
      LATENT_DIM_1 = 64
      DIM_PIX_2    = 512
      DIM_4        = 512
      LATENT_DIM_2 = 512
  
      ALPHA1_ITERS = 2000
      ALPHA2_ITERS = 5000
      KL_PENALTY = 1.00
      BETA_ITERS = 1000
  
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 1000,
          'stop_after': 200000,
          'callback_every': 20000
      }
  
      LR = 1e-3
  
      LR_DECAY_AFTER = 180000
      LR_DECAY_FACTOR = 1e-1
  
      BATCH_SIZE = 64
      N_CHANNELS = 3
      HEIGHT = 32
      WIDTH = 32
  
      LATENTS1_HEIGHT = 8
      LATENTS1_WIDTH = 8
  
  elif SETTINGS == '32px_big':
  
      MODE = 'two_level'
  
      EMBED_INPUTS = False
  
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 256
      DIM_1        = 128
      DIM_2        = 256
      DIM_3        = 512
      LATENT_DIM_1 = 128
      DIM_PIX_2    = 512
      DIM_4        = 512
      LATENT_DIM_2 = 512
  
      ALPHA1_ITERS = 2000
      ALPHA2_ITERS = 5000
      KL_PENALTY = 1.00
      BETA_ITERS = 1000
  
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 1000,
          'stop_after': 300000,
          'callback_every': 20000
      }
  
      VANILLA = False
      LR = 1e-3
  
      LR_DECAY_AFTER = 300000
      LR_DECAY_FACTOR = 1e-1
  
      BATCH_SIZE = 64
      N_CHANNELS = 3
      HEIGHT = 32
      WIDTH = 32
      LATENTS1_HEIGHT = 8
      LATENTS1_WIDTH = 8
  
  elif SETTINGS == '64px_small':
      MODE = 'two_level'
  
      EMBED_INPUTS = True
  
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 128
      DIM_0        = 64
      DIM_1        = 64
      DIM_2        = 128
      LATENT_DIM_1 = 64
      DIM_PIX_2    = 256
      DIM_3        = 256
      DIM_4        = 512
      LATENT_DIM_2 = 512
  
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 10000,
          'stop_after': 200000,
          'callback_every': 50000
      }
  
      VANILLA = False
      LR = 1e-3
  
      LR_DECAY_AFTER = 180000
      LR_DECAY_FACTOR = .1
  
      ALPHA1_ITERS = 2000
      ALPHA2_ITERS = 10000
      KL_PENALTY = 1.0
      BETA_ITERS = 1000
  
      BATCH_SIZE = 64
      N_CHANNELS = 3
      HEIGHT = 64
      WIDTH = 64
      LATENTS1_WIDTH = 16
      LATENTS1_HEIGHT = 16
  
  elif SETTINGS == '64px_big':
      MODE = 'two_level'
  
      EMBED_INPUTS = True
  
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 384
      DIM_0        = 192
      DIM_1        = 256
      DIM_2        = 512
      LATENT_DIM_1 = 64
      DIM_PIX_2    = 512
      DIM_3        = 512
      DIM_4        = 512
      LATENT_DIM_2 = 512
  
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 10000,
          'stop_after': 400000,
          'callback_every': 50000
      }
  
      VANILLA = False
      LR = 1e-3
  
      LR_DECAY_AFTER = 180000
      LR_DECAY_FACTOR = .5
  
      ALPHA1_ITERS = 1000
      ALPHA2_ITERS = 10000
      KL_PENALTY = 1.00
      BETA_ITERS = 500
  
      BATCH_SIZE = 48
      N_CHANNELS = 3
      HEIGHT = 64
      WIDTH = 64
      LATENTS1_WIDTH = 16
      LATENTS1_HEIGHT = 16
  
  elif SETTINGS=='64px_big_onelevel':
  
      # two_level uses Enc1/Dec1 for the bottom level, Enc2/Dec2 for the top level
      # one_level uses EncFull/DecFull for the bottom (and only) level
      MODE = 'one_level'
  
      # Whether to treat pixel inputs to the model as real-valued (as in the 
      # original PixelCNN) or discrete (gets better likelihoods).
      EMBED_INPUTS = True
  
      # Turn on/off the bottom-level PixelCNN in Dec1/DecFull
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 384
      DIM_0        = 192
      DIM_1        = 256
      DIM_2        = 512
      DIM_3        = 512
      DIM_4        = 512
      LATENT_DIM_2 = 512
  
      ALPHA1_ITERS = 50000
      ALPHA2_ITERS = 50000
      KL_PENALTY = 1.0
      BETA_ITERS = 1000
  
      # In Dec2, we break each spatial location into N blocks (analogous to channels
      # in the original PixelCNN) and model each spatial location autoregressively
      # as P(x)=P(x0)*P(x1|x0)*P(x2|x0,x1)... In my experiments values of N > 1
      # actually hurt performance. Unsure why; might be a bug.
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 10000,
          'stop_after': 400000,
          'callback_every': 50000
      }
      LR = 1e-3
  
      LR_DECAY_AFTER = 180000
      LR_DECAY_FACTOR = 0.5
  
      BATCH_SIZE = 48
      N_CHANNELS = 3
      HEIGHT = 64
      WIDTH = 64
  
      # These aren't actually used for one-level models but some parts
      # of the code still depend on them being defined.
      LATENT_DIM_1 = 64
      LATENTS1_HEIGHT = 7
      LATENTS1_WIDTH = 7
  
  elif SETTINGS=='32px_cifar':
  
      # two_level uses Enc1/Dec1 for the bottom level, Enc2/Dec2 for the top level
      # one_level uses EncFull/DecFull for the bottom (and only) level
      MODE = 'one_level'
  
      # Whether to treat pixel inputs to the model as real-valued (as in the 
      # original PixelCNN) or discrete (gets better likelihoods).
      EMBED_INPUTS = True
  
      # Turn on/off the bottom-level PixelCNN in Dec1/DecFull
      PIXEL_LEVEL_PIXCNN = True
      HIGHER_LEVEL_PIXCNN = True
  
      DIM_EMBED    = 16
      DIM_PIX_1    = 192 #LEILA EDIT: was previously 384
      DIM_0        = 96 #LEILA EDIT: was previously 192
      DIM_1        = 128 #LEILA EDIT: was previously 256
      DIM_2        = 256 #LEILA EDIT: was previously 512
      DIM_3        = 256 #LEILA EDIT: was previously 512
      DIM_4        = 256 #LEILA EDIT: was previously 512
      LATENT_DIM_2 = 2 #LEILA EDIT2: was previously 256
  
      ALPHA1_ITERS = 50000
      ALPHA2_ITERS = 50000
      KL_PENALTY = 1.0
      BETA_ITERS = 1000
  
      # In Dec2, we break each spatial location into N blocks (analogous to channels
      # in the original PixelCNN) and model each spatial location autoregressively
      # as P(x)=P(x0)*P(x1|x0)*P(x2|x0,x1)... In my experiments values of N > 1
      # actually hurt performance. Unsure why; might be a bug.
      PIX_2_N_BLOCKS = 1
  
      TIMES = {
          'test_every': 10000,
          'stop_after': 400000,
          'callback_every': 50000
      }
      
      LR = 1e-3
  
      LR_DECAY_AFTER = 180000
      LR_DECAY_FACTOR = 0.5
  
      BATCH_SIZE = 50 # 48
      N_CHANNELS = 3
      HEIGHT = 32 #64
      WIDTH = 32 #64
  
      # These aren't actually used for one-level models but some parts
      # of the code still depend on them being defined.
      LATENT_DIM_1 = 32 #LEILAEDIT: was previously 64
      LATENTS1_HEIGHT = 7
      LATENTS1_WIDTH = 7
      
  if DATASET == 'mnist_256':
      train_data, dev_data, test_data = lib.mnist_256.load(BATCH_SIZE, BATCH_SIZE)
  elif DATASET == 'lsun_32':
      train_data, dev_data = lib.lsun_bedrooms.load(BATCH_SIZE, downsample=True)
  elif DATASET == 'lsun_64':
      train_data, dev_data = lib.lsun_bedrooms.load(BATCH_SIZE, downsample=False)
  elif DATASET == 'imagenet_64':
      train_data, dev_data = lib.small_imagenet.load(BATCH_SIZE)
  elif DATASET == 'cifar10':
      train_data, dev_data, test_data = lib.cifar_256.load(BATCH_SIZE) #LEILAEDIT
  
  lib.print_model_settings(locals().copy())
  
  DEVICES = ['/gpu:{}'.format(i) for i in xrange(N_GPUS)]
  
  lib.ops.conv2d.enable_default_weightnorm()
  lib.ops.linear.enable_default_weightnorm()
  
  with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as session:
      bn_is_training = tf.placeholder(tf.bool, shape=None, name='bn_is_training')
      bn_stats_iter = tf.placeholder(tf.int32, shape=None, name='bn_stats_iter')
      total_iters = tf.placeholder(tf.int32, shape=None, name='total_iters')
      all_images = tf.placeholder(tf.int32, shape=[None, N_CHANNELS, HEIGHT, WIDTH], name='all_images')
      all_latents1 = tf.placeholder(tf.float32, shape=[None, LATENT_DIM_1, LATENTS1_HEIGHT, LATENTS1_WIDTH], name='all_latents1')
  
      split_images = tf.split(all_images, len(DEVICES), axis=0)
      split_latents1 = tf.split(all_images, len(DEVICES), axis=0)
  
      tower_cost = []
      tower_outputs1_sample = []
  
      for device_index, (device, images, latents1_sample) in enumerate(zip(DEVICES, split_images, split_latents1)):
          with tf.device(device):
  
              def nonlinearity(x):
                  return tf.nn.elu(x)
  
              def pixcnn_gated_nonlinearity(a, b):
                  return tf.sigmoid(a) * tf.tanh(b)
  
              def SubpixelConv2D(*args, **kwargs):
                  kwargs['output_dim'] = 4*kwargs['output_dim']
                  output = lib.ops.conv2d.Conv2D(*args, **kwargs)
                  output = tf.transpose(output, [0,2,3,1])
                  output = tf.depth_to_space(output, 2)
                  output = tf.transpose(output, [0,3,1,2])
                  return output
  
              def ResidualBlock(name, input_dim, output_dim, inputs, filter_size, mask_type=None, resample=None, he_init=True):
                  """
                  resample: None, 'down', or 'up'
                  """
                  if mask_type != None and resample != None:
                      raise Exception('Unsupported configuration')
  
                  if resample=='down':
                      conv_shortcut = functools.partial(lib.ops.conv2d.Conv2D, stride=2)
                      conv_1        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim, output_dim=input_dim)
                      conv_2        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim, output_dim=output_dim, stride=2)
                  elif resample=='up':
                      conv_shortcut = SubpixelConv2D
                      conv_1        = functools.partial(SubpixelConv2D, input_dim=input_dim, output_dim=output_dim)
                      conv_2        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=output_dim, output_dim=output_dim)
                  elif resample==None:
                      conv_shortcut = lib.ops.conv2d.Conv2D
                      conv_1        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim, output_dim=output_dim)
                      conv_2        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=output_dim, output_dim=output_dim)
                  else:
                      raise Exception('invalid resample value')
  
                  if output_dim==input_dim and resample==None:
                      shortcut = inputs # Identity skip-connection
                  else:
                      shortcut = conv_shortcut(name+'.Shortcut', input_dim=input_dim, output_dim=output_dim, filter_size=1, mask_type=mask_type, he_init=False, biases=True, inputs=inputs)
  
                  output = inputs
                  if mask_type == None:
                      output = nonlinearity(output)
                      output = conv_1(name+'.Conv1', filter_size=filter_size, mask_type=mask_type, inputs=output, he_init=he_init, weightnorm=False)
                      output = nonlinearity(output)
                      output = conv_2(name+'.Conv2', filter_size=filter_size, mask_type=mask_type, inputs=output, he_init=he_init, weightnorm=False, biases=False)
                      if device_index == 0:
                          output = lib.ops.batchnorm.Batchnorm(name+'.BN', [0,2,3], output, bn_is_training, bn_stats_iter)
                      else:
                          output = lib.ops.batchnorm.Batchnorm(name+'.BN', [0,2,3], output, bn_is_training, bn_stats_iter, update_moving_stats=False)
                  else:
                      output = nonlinearity(output)
                      output_a = conv_1(name+'.Conv1A', filter_size=filter_size, mask_type=mask_type, inputs=output, he_init=he_init)
                      output_b = conv_1(name+'.Conv1B', filter_size=filter_size, mask_type=mask_type, inputs=output, he_init=he_init)
                      output = pixcnn_gated_nonlinearity(output_a, output_b)
                      output = conv_2(name+'.Conv2', filter_size=filter_size, mask_type=mask_type, inputs=output, he_init=he_init)
  
                  return shortcut + output
  
              
              # Only for 32px_cifar, 64px_big_onelevel, and MNIST. Needs modification for others.
              def EncFull(images):
                  output = images
  
                  if WIDTH == 32: #64 
                      if EMBED_INPUTS:
                          output = lib.ops.conv2d.Conv2D('EncFull.Input', input_dim=N_CHANNELS*DIM_EMBED, output_dim=DIM_0, filter_size=1, inputs=output, he_init=False)
                      else:
                          output = lib.ops.conv2d.Conv2D('EncFull.Input', input_dim=N_CHANNELS, output_dim=DIM_0, filter_size=1, inputs=output, he_init=False)
  
                      output = ResidualBlock('EncFull.Res1', input_dim=DIM_0, output_dim=DIM_0, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res2', input_dim=DIM_0, output_dim=DIM_1, filter_size=3, resample='down', inputs=output)
                      output = ResidualBlock('EncFull.Res3', input_dim=DIM_1, output_dim=DIM_1, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res4', input_dim=DIM_1, output_dim=DIM_1, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res5', input_dim=DIM_1, output_dim=DIM_2, filter_size=3, resample='down', inputs=output)
                      output = ResidualBlock('EncFull.Res6', input_dim=DIM_2, output_dim=DIM_2, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res7', input_dim=DIM_2, output_dim=DIM_2, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res8', input_dim=DIM_2, output_dim=DIM_3, filter_size=3, resample='down', inputs=output)
                      output = ResidualBlock('EncFull.Res9', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res10', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res11', input_dim=DIM_3, output_dim=DIM_4, filter_size=3, resample='down', inputs=output)
                      output = ResidualBlock('EncFull.Res12', input_dim=DIM_4, output_dim=DIM_4, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res13', input_dim=DIM_4, output_dim=DIM_4, filter_size=3, resample=None, inputs=output)
                      output = tf.reshape(output, [-1, 2*2*DIM_4])
                      output = lib.ops.linear.Linear('EncFull.Output', input_dim=2*2*DIM_4, output_dim=2*LATENT_DIM_2, initialization='glorot', inputs=output)
                  else:
                      if EMBED_INPUTS:
                          output = lib.ops.conv2d.Conv2D('EncFull.Input', input_dim=N_CHANNELS*DIM_EMBED, output_dim=DIM_1, filter_size=1, inputs=output, he_init=False)
                      else:
                          output = lib.ops.conv2d.Conv2D('EncFull.Input', input_dim=N_CHANNELS, output_dim=DIM_1, filter_size=1, inputs=output, he_init=False)
  
                      output = ResidualBlock('EncFull.Res1', input_dim=DIM_1, output_dim=DIM_1, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res2', input_dim=DIM_1, output_dim=DIM_2, filter_size=3, resample='down', inputs=output)
                      output = ResidualBlock('EncFull.Res3', input_dim=DIM_2, output_dim=DIM_2, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res4', input_dim=DIM_2, output_dim=DIM_3, filter_size=3, resample='down', inputs=output)
                      output = ResidualBlock('EncFull.Res5', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, inputs=output)
                      output = ResidualBlock('EncFull.Res6', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, inputs=output)
                      output = tf.reduce_mean(output, reduction_indices=[2,3])
                      output = lib.ops.linear.Linear('EncFull.Output', input_dim=DIM_3, output_dim=2*LATENT_DIM_2, initialization='glorot', inputs=output)
  
                  return output
  
              # Only for 32px_CIFAR, 64px_big_onelevel and MNIST. Needs modification for others.
              def DecFull(latents, images):
                  output = tf.clip_by_value(latents, -50., 50.)
  
                  if WIDTH == 32: # 64:LEILAEDIT. Also changed 4*4 to 2*2 and 4,4 to 2,2 in the two lines below
                      output = lib.ops.linear.Linear('DecFull.Input', input_dim=LATENT_DIM_2, output_dim=2*2*DIM_4, initialization='glorot', inputs=output)
                      output = tf.reshape(output, [-1, DIM_4, 2, 2])
                      output = ResidualBlock('DecFull.Res2', input_dim=DIM_4, output_dim=DIM_4, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res3', input_dim=DIM_4, output_dim=DIM_4, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res4', input_dim=DIM_4, output_dim=DIM_3, filter_size=3, resample='up', he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res5', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res6', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res7', input_dim=DIM_3, output_dim=DIM_2, filter_size=3, resample='up', he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res8', input_dim=DIM_2, output_dim=DIM_2, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res9', input_dim=DIM_2, output_dim=DIM_2, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res10', input_dim=DIM_2, output_dim=DIM_1, filter_size=3, resample='up', he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res11', input_dim=DIM_1, output_dim=DIM_1, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res12', input_dim=DIM_1, output_dim=DIM_1, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res13', input_dim=DIM_1, output_dim=DIM_0, filter_size=3, resample='up', he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res14', input_dim=DIM_0, output_dim=DIM_0, filter_size=3, resample=None, he_init=True, inputs=output)
                  else:
                      output = lib.ops.linear.Linear('DecFull.Input', input_dim=LATENT_DIM_2, output_dim=DIM_3, initialization='glorot', inputs=output)
                      output = tf.reshape(tf.tile(tf.reshape(output, [-1, DIM_3, 1]), [1, 1, 49]), [-1, DIM_3, 7, 7])
                      output = ResidualBlock('DecFull.Res2', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res3', input_dim=DIM_3, output_dim=DIM_3, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res4', input_dim=DIM_3, output_dim=DIM_2, filter_size=3, resample='up', he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res5', input_dim=DIM_2, output_dim=DIM_2, filter_size=3, resample=None, he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res6', input_dim=DIM_2, output_dim=DIM_1, filter_size=3, resample='up', he_init=True, inputs=output)
                      output = ResidualBlock('DecFull.Res7', input_dim=DIM_1, output_dim=DIM_1, filter_size=3, resample=None, he_init=True, inputs=output)
  
                  if WIDTH == 32: #64:
                      dim = DIM_0
                  else:
                      dim = DIM_1
  
                  if PIXEL_LEVEL_PIXCNN: #LEILAEDIT
  
                      if EMBED_INPUTS:
                          masked_images = lib.ops.conv2d.Conv2D('DecFull.Pix1', input_dim=N_CHANNELS*DIM_EMBED, output_dim=dim, filter_size=3, inputs=images, mask_type=('a', N_CHANNELS), he_init=False)
                      else:
                          masked_images = lib.ops.conv2d.Conv2D('DecFull.Pix1', input_dim=N_CHANNELS, output_dim=dim, filter_size=3, inputs=images, mask_type=('a', N_CHANNELS), he_init=False)
  
                      # Warning! Because of the masked convolutions it's very important that masked_images comes first in this concat
                      output = tf.concat([masked_images, output], axis=1)
  
                      output = ResidualBlock('DecFull.Pix2Res', input_dim=2*dim,   output_dim=DIM_PIX_1, filter_size=3, mask_type=('b', N_CHANNELS), inputs=output)
                      output = ResidualBlock('DecFull.Pix3Res', input_dim=DIM_PIX_1, output_dim=DIM_PIX_1, filter_size=3, mask_type=('b', N_CHANNELS), inputs=output)
                      output = ResidualBlock('DecFull.Pix4Res', input_dim=DIM_PIX_1, output_dim=DIM_PIX_1, filter_size=3, mask_type=('b', N_CHANNELS), inputs=output)
                      if WIDTH != 32: #64: LEILAEDIT
                          output = ResidualBlock('DecFull.Pix5Res', input_dim=DIM_PIX_1, output_dim=DIM_PIX_1, filter_size=3, mask_type=('b', N_CHANNELS), inputs=output)
  
                      output = lib.ops.conv2d.Conv2D('Dec1.Out', input_dim=DIM_PIX_1, output_dim=256*N_CHANNELS, filter_size=1, mask_type=('b', N_CHANNELS), he_init=False, inputs=output)
                    
                  else:
  
                      output = lib.ops.conv2d.Conv2D('Dec1.Out', input_dim=dim, output_dim=256*N_CHANNELS, filter_size=1, he_init=False, inputs=output)
  
                  return tf.transpose(
                      tf.reshape(output, [-1, 256, N_CHANNELS, HEIGHT, WIDTH]),
                      [0,2,3,4,1]
                  )
  
              def split(mu_and_logsig):
                  mu, logsig = tf.split(mu_and_logsig, 2, axis=1)
                  sig = 0.5 * (tf.nn.softsign(logsig)+1)
                  logsig = tf.log(sig)
                  return mu, logsig, sig
           
              def clamp_logsig_and_sig(logsig, sig):
                  # Early during training (see BETA_ITERS), stop sigma from going too low
                  floor = 1. - tf.minimum(1., tf.cast(total_iters, 'float32') / BETA_ITERS)
                  log_floor = tf.log(floor)
                  return tf.maximum(logsig, log_floor), tf.maximum(sig, floor)
  
  
              scaled_images = (tf.cast(images, 'float32') - 128.) / 64.
              if EMBED_INPUTS:
                  embedded_images = lib.ops.embedding.Embedding('Embedding', 256, DIM_EMBED, images)
                  embedded_images = tf.transpose(embedded_images, [0,4,1,2,3])
                  embedded_images = tf.reshape(embedded_images, [-1, DIM_EMBED*N_CHANNELS, HEIGHT, WIDTH])
  
              if MODE == 'one_level':
  
                  # Layer 1
  
                  if EMBED_INPUTS:
                      mu_and_logsig1 = EncFull(embedded_images)
                  else:
                      mu_and_logsig1 = EncFull(scaled_images)
                  mu1, logsig1, sig1 = split(mu_and_logsig1)
  
                  eps = tf.random_normal(tf.shape(mu1))
                  latents1 = mu1 + (eps * sig1)
  
                  if EMBED_INPUTS:
                      outputs1 = DecFull(latents1, embedded_images)
                  else:
                      outputs1 = DecFull(latents1, scaled_images)
  
                  reconst_cost = tf.reduce_mean(
                      tf.nn.sparse_softmax_cross_entropy_with_logits(
                          logits=tf.reshape(outputs1, [-1, 256]),
                          labels=tf.reshape(images, [-1])
                      )
                  )
  
                  # Assembly
  
                  # An alpha of exactly 0 can sometimes cause inf/nan values, so we're
                  # careful to avoid it.
                  alpha = tf.minimum(1., tf.cast(total_iters+1, 'float32') / ALPHA1_ITERS) * KL_PENALTY
  
                  kl_cost_1 = tf.reduce_mean(
                      lib.ops.kl_unit_gaussian.kl_unit_gaussian(
                          mu1, 
                          logsig1,
                          sig1
                      )
                  )
  
                  kl_cost_1 *= float(LATENT_DIM_2) / (N_CHANNELS * WIDTH * HEIGHT)
  
                  cost = reconst_cost + (alpha * kl_cost_1)
  
              tower_cost.append(cost)
  
      full_cost = tf.reduce_mean(
          tf.concat([tf.expand_dims(x, 0) for x in tower_cost], axis=0), 0
      )
  
      # Sampling
  
      if MODE == 'one_level':
  
          ch_sym = tf.placeholder(tf.int32, shape=None)
          y_sym = tf.placeholder(tf.int32, shape=None)
          x_sym = tf.placeholder(tf.int32, shape=None)
          logits = tf.reshape(tf.slice(outputs1, tf.stack([0, ch_sym, y_sym, x_sym, 0]), tf.stack([-1, 1, 1, 1, -1])), [-1, 256])
          dec1_fn_out = tf.multinomial(logits, 1)[:, 0]
          def dec1_fn(_latents, _targets, _ch, _y, _x):
              return session.run(dec1_fn_out, feed_dict={latents1: _latents, images: _targets, ch_sym: _ch, y_sym: _y, x_sym: _x, total_iters: 99999, bn_is_training: False, bn_stats_iter:0})
  
          def enc_fn(_images):
              return session.run(latents1, feed_dict={images: _images, total_iters: 99999, bn_is_training: False, bn_stats_iter:0})
  
          #sample_fn_latents1 = np.random.normal(size=(1, LATENT_DIM_2)).astype('float32') # changed 8 to 1
           
          def generate_and_save_samples(tag):
              #def color_grid_vis(X, nh, nw, save_path):
              #    # from github.com/Newmu
              #    X = X.transpose(0,2,3,1)
              #    h, w = X[0].shape[:2]
              #    img = np.zeros((h*nh, w*nw, 3))
              #    for n, x in enumerate(X):
              #        j = n/nw
              #        i = n%nw
              #        img[j*h:j*h+h, i*w:i*w+w, :] = x
              #    imsave(save_path, img)
  
              latents1_copied = np.zeros((1, LATENT_DIM_2), dtype='float32') # changed 8 to 1
              #for i in xrange(1): # changed 8 to 1
              #    latents1_copied[i::1] = sample_fn_latents1 # changed 8 to 1
  
              #samples = np.zeros(
              #    (1, N_CHANNELS, HEIGHT, WIDTH), # changed 64 to 1
              #    dtype='int32'
              #)
  
              #print "Generating samples"
              #for y in xrange(HEIGHT):
              #    for x in xrange(WIDTH):
              #        for ch in xrange(N_CHANNELS):
              #            next_sample = dec1_fn(latents1_copied, samples, ch, y, x)
              #            samples[:,ch,y,x] = next_sample
  
              #print "Saving samples"
              #color_grid_vis(
              #    samples, 
              #    1, 
              #    1, 
              #    'samples_filter_3_{}.png'.format(tag) # changed to 1 and 1
              #)
  
      # Train!
  
      if MODE == 'one_level':
          prints=[
              ('alpha', alpha), 
              ('reconst', reconst_cost), 
              ('kl1', kl_cost_1)
          ]
        
      decayed_lr = tf.train.exponential_decay(
          LR,
          total_iters,
          LR_DECAY_AFTER,
          LR_DECAY_FACTOR,
          staircase=True
      )
     
      lib.train_loop_cifar_filter_3_dim_2.train_loop(
          session=session,
          inputs=[total_iters, all_images],
          inject_iteration=True,
          bn_vars=(bn_is_training, bn_stats_iter),
          cost=full_cost,
          stop_after=TIMES['stop_after'],
          prints=prints,
          optimizer=tf.train.AdamOptimizer(decayed_lr),
          train_data=train_data,
          test_data=dev_data,
          callback=generate_and_save_samples,
          callback_every=TIMES['callback_every'],
          test_every=TIMES['test_every'],
          save_checkpoints=True
      )
  