# Copyright 2016 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Model architecture for predictive model, including CDNA, DNA, and STP."""

import numpy as np
import tensorflow as tf

import tensorflow.contrib.slim as slim
from tensorflow.python.platform import flags
from tensorflow.contrib.layers.python import layers as tf_layers
from lstm_ops import basic_conv_lstm_cell

FLAGS = flags.FLAGS

# Amount to use when lower bounding tensors
RELU_SHIFT = 1e-12

# kernel size for DNA and CDNA.
DNA_KERN_SIZE = 5


def kl_divergence(mu, log_sigma):
  """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).

  Args:
    mu: mu parameter of the distribution.
    log_sigma: log(sigma) parameter of the distribution.
  Returns:
    the KL loss.
  """

  return -.5 * tf.reduce_sum(1. + log_sigma - tf.square(mu) - tf.exp(log_sigma),
                             axis=1)


def construct_latent_tower(images):
  """Builds convolutional latent tower for stochastic model.

  At training time this tower generates a latent distribution (mean and std)
  conditioned on the entire video. This latent variable will be fed to the
  main tower as an extra variable to be used for future frames prediction.
  At inference time, the tower is disabled and only returns latents sampled
  from N(0,1).
  If the multi_latent flag is on, a different latent for every timestep would
  be generated.

  Args:
    images: tensor of ground truth image sequences
  Returns:
    latent_mean: predicted latent mean
    latent_std: predicted latent standard deviation
    latent_loss: loss of the latent twoer
    samples: random samples sampled from standard guassian
  """

  with slim.arg_scope([slim.conv2d], reuse=False):
    stacked_images = tf.concat(images, 3)

    latent_enc1 = slim.conv2d(
        stacked_images,
        32, [3, 3],
        stride=2,
        scope='latent_conv1',
        normalizer_fn=tf_layers.layer_norm,
        normalizer_params={'scope': 'latent_norm1'})

    latent_enc2 = slim.conv2d(
        latent_enc1,
        64, [3, 3],
        stride=2,
        scope='latent_conv2',
        normalizer_fn=tf_layers.layer_norm,
        normalizer_params={'scope': 'latent_norm2'})

    latent_enc3 = slim.conv2d(
        latent_enc2,
        64, [3, 3],
        stride=1,
        scope='latent_conv3',
        normalizer_fn=tf_layers.layer_norm,
        normalizer_params={'scope': 'latent_norm3'})

    latent_mean = slim.conv2d(
        latent_enc3,
        FLAGS.latent_channels, [3, 3],
        stride=2,
        activation_fn=None,
        scope='latent_mean',
        normalizer_fn=tf_layers.layer_norm,
        normalizer_params={'scope': 'latent_norm_mean'})

    latent_std = slim.conv2d(
        latent_enc3,
        FLAGS.latent_channels, [3, 3],
        stride=2,
        scope='latent_std',
        normalizer_fn=tf_layers.layer_norm,
        normalizer_params={'scope': 'latent_std_norm'})

    latent_std += FLAGS.latent_std_min

    divergence = kl_divergence(latent_mean, latent_std)
    latent_loss = tf.reduce_mean(divergence)

  if FLAGS.multi_latent:
    # timestep x batch_size x latent_size
    samples = tf.random_normal(
        [FLAGS.sequence_length-1] + latent_mean.shape, 0, 1,
        dtype=tf.float32)
  else:
    # batch_size x latent_size
    samples = tf.random_normal(latent_mean.shape, 0, 1, dtype=tf.float32)

  if FLAGS.inference_time:
    # No latent tower at inference time, just standard gaussian.
    return None, None, None, samples
  else:
    return latent_mean, latent_std, latent_loss, samples


def construct_model(images,
                    actions=None,
                    states=None,
                    iter_num=-1.0,
                    k=-1,
                    use_state=True,
                    num_masks=10,
                    stp=False,
                    cdna=True,
                    dna=False,
                    context_frames=2):
  """Build convolutional lstm video predictor using STP, CDNA, or DNA.

  Args:
    images: tensor of ground truth image sequences
    actions: tensor of action sequences
    states: tensor of ground truth state sequences
    iter_num: tensor of the current training iteration (for sched. sampling)
    k: constant used for scheduled sampling. -1 to feed in own prediction.
    use_state: True to include state and action in prediction
    num_masks: the number of different pixel motion predictions (and
               the number of masks for each of those predictions)
    stp: True to use Spatial Transformer Predictor (STP)
    cdna: True to use Convoluational Dynamic Neural Advection (CDNA)
    dna: True to use Dynamic Neural Advection (DNA)
    context_frames: number of ground truth frames to pass in before
                    feeding in own predictions
  Returns:
    gen_images: predicted future image frames
    gen_states: predicted future states

  Raises:
    ValueError: if more than one network option specified or more than 1 mask
    specified for DNA model.
  """
  # Each image is being used twice, in latent tower and main tower.
  # This is to make sure we are using the *same* image for both, ...
  # ... given how TF queues work. 
  images = [tf.identity(image) for image in images]

  if stp + cdna + dna != 1:
    raise ValueError('More than one, or no network option specified.')
  batch_size, img_height, img_width, color_channels = images[0].get_shape()[0:4]
  lstm_func = basic_conv_lstm_cell

  # Generated robot states and images.
  gen_states, gen_images = [], []
  current_state = states[0]

  if k == -1:
    feedself = True
  else:
    # Scheduled sampling:
    # Calculate number of ground-truth frames to pass in.
    num_ground_truth = tf.to_int32(
        tf.round(tf.to_float(batch_size) * (k / (k + tf.exp(iter_num / k)))))
    feedself = False

  # LSTM state sizes and states.
  lstm_size = np.int32(np.array([32, 32, 64, 64, 128, 64, 32]))
  lstm_state1, lstm_state2, lstm_state3, lstm_state4 = None, None, None, None
  lstm_state5, lstm_state6, lstm_state7 = None, None, None

  # Latent tower
  latent_loss = 0.0
  if FLAGS.stochastic_model:
    latent_tower_outputs = construct_latent_tower(images)
    latent_mean, latent_std, latent_loss, samples = latent_tower_outputs

  # Main tower    
  for image, action in zip(images[:-1], actions[:-1]):
    # Reuse variables after the first timestep.
    reuse = bool(gen_images)

    done_warm_start = len(gen_images) > context_frames - 1
    with slim.arg_scope(
        [lstm_func, slim.layers.conv2d, slim.layers.fully_connected,
         tf_layers.layer_norm, slim.layers.conv2d_transpose],
        reuse=reuse):

      if feedself and done_warm_start:
        # Feed in generated image.
        prev_image = gen_images[-1]
      elif done_warm_start:
        # Scheduled sampling
        prev_image = scheduled_sample(image, gen_images[-1], batch_size,
                                      num_ground_truth)
      else:
        # Always feed in ground_truth
        prev_image = image

      # Predicted state is always fed back in
      state_action = tf.concat(axis=1, values=[action, current_state])

      enc0 = slim.layers.conv2d(
          prev_image,
          32, [5, 5],
          stride=2,
          scope='scale1_conv1',
          normalizer_fn=tf_layers.layer_norm,
          normalizer_params={'scope': 'layer_norm1'})

      hidden1, lstm_state1 = lstm_func(
          enc0, lstm_state1, lstm_size[0], scope='state1')
      hidden1 = tf_layers.layer_norm(hidden1, scope='layer_norm2')
      hidden2, lstm_state2 = lstm_func(
          hidden1, lstm_state2, lstm_size[1], scope='state2')
      hidden2 = tf_layers.layer_norm(hidden2, scope='layer_norm3')
      enc1 = slim.layers.conv2d(
          hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope='conv2')

      hidden3, lstm_state3 = lstm_func(
          enc1, lstm_state3, lstm_size[2], scope='state3')
      hidden3 = tf_layers.layer_norm(hidden3, scope='layer_norm4')
      hidden4, lstm_state4 = lstm_func(
          hidden3, lstm_state4, lstm_size[3], scope='state4')
      hidden4 = tf_layers.layer_norm(hidden4, scope='layer_norm5')
      enc2 = slim.layers.conv2d(
          hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope='conv3')

      # Pass in state and action.
      smear = tf.reshape(
          state_action,
          [int(batch_size), 1, 1, int(state_action.get_shape()[1])])
      smear = tf.tile(
          smear, [1, int(enc2.get_shape()[1]), int(enc2.get_shape()[2]), 1])
      if use_state:
        enc2 = tf.concat(axis=3, values=[enc2, smear])
        
      # Setup latent
      if FLAGS.stochastic_model:
        latent = samples
        if FLAGS.multi_latent:
          latent = samples[timestep]
        if not FLAGS.inference_time:
          latent = tf.cond(iter_num < FLAGS.num_iterations_1st_stage,
                           lambda: tf.identity(latent),
                           lambda: latent_mean + tf.exp(latent_std / 2.0) * latent)
        with tf.control_dependencies([latent]):
          enc2 = tf.concat([enc2, latent], 3)
        
      enc3 = slim.layers.conv2d(
          enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope='conv4')

      hidden5, lstm_state5 = lstm_func(
          enc3, lstm_state5, lstm_size[4], scope='state5')  # last 8x8
      hidden5 = tf_layers.layer_norm(hidden5, scope='layer_norm6')
      enc4 = slim.layers.conv2d_transpose(
          hidden5, hidden5.get_shape()[3], 3, stride=2, scope='convt1')

      hidden6, lstm_state6 = lstm_func(
          enc4, lstm_state6, lstm_size[5], scope='state6')  # 16x16
      hidden6 = tf_layers.layer_norm(hidden6, scope='layer_norm7')
      # Skip connection.
      hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16

      enc5 = slim.layers.conv2d_transpose(
          hidden6, hidden6.get_shape()[3], 3, stride=2, scope='convt2')
      hidden7, lstm_state7 = lstm_func(
          enc5, lstm_state7, lstm_size[6], scope='state7')  # 32x32
      hidden7 = tf_layers.layer_norm(hidden7, scope='layer_norm8')

      # Skip connection.
      hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32

      enc6 = slim.layers.conv2d_transpose(
          hidden7,
          hidden7.get_shape()[3], 3, stride=2, scope='convt3', activation_fn=None,
          normalizer_fn=tf_layers.layer_norm,
          normalizer_params={'scope': 'layer_norm9'})

      if dna:
        # Using largest hidden state for predicting untied conv kernels.
        enc7 = slim.layers.conv2d_transpose(
            enc6, DNA_KERN_SIZE**2, 1, stride=1, scope='convt4', activation_fn=None)
      else:
        # Using largest hidden state for predicting a new image layer.
        enc7 = slim.layers.conv2d_transpose(
            enc6, color_channels, 1, stride=1, scope='convt4', activation_fn=None)
        # This allows the network to also generate one image from scratch,
        # which is useful when regions of the image become unoccluded.
        transformed = [tf.nn.sigmoid(enc7)]

      if stp:
        stp_input0 = tf.reshape(hidden5, [int(batch_size), -1])
        stp_input1 = slim.layers.fully_connected(
            stp_input0, 100, scope='fc_stp')
        transformed += stp_transformation(prev_image, stp_input1, num_masks)
      elif cdna:
        cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
        transformed += cdna_transformation(prev_image, cdna_input, num_masks,
                                           int(color_channels))
      elif dna:
        # Only one mask is supported (more should be unnecessary).
        if num_masks != 1:
          raise ValueError('Only one mask is supported for DNA model.')
        transformed = [dna_transformation(prev_image, enc7)]

      masks = slim.layers.conv2d_transpose(
          enc6, num_masks + 1, 1, stride=1, scope='convt7', activation_fn=None)
      masks = tf.reshape(
          tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
          [int(batch_size), int(img_height), int(img_width), num_masks + 1])
      mask_list = tf.split(axis=3, num_or_size_splits=num_masks + 1, value=masks)
      output = mask_list[0] * prev_image
      for layer, mask in zip(transformed, mask_list[1:]):
        output += layer * mask
      gen_images.append(output)

      current_state = slim.layers.fully_connected(
          state_action,
          int(current_state.get_shape()[1]),
          scope='state_pred',
          activation_fn=None)
      gen_states.append(current_state)

  return gen_images, gen_states, latent_loss


## Utility functions
def stp_transformation(prev_image, stp_input, num_masks):
  """Apply spatial transformer predictor (STP) to previous image.

  Args:
    prev_image: previous image to be transformed.
    stp_input: hidden layer to be used for computing STN parameters.
    num_masks: number of masks and hence the number of STP transformations.
  Returns:
    List of images transformed by the predicted STP parameters.
  """
  # Only import spatial transformer if needed.
  from spatial_transformer import transformer

  identity_params = tf.convert_to_tensor(
      np.array([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], np.float32))
  transformed = []
  for i in range(num_masks - 1):
    params = slim.layers.fully_connected(
        stp_input, 6, scope='stp_params' + str(i),
        activation_fn=None) + identity_params
    transformed.append(transformer(prev_image, params))

  return transformed


def cdna_transformation(prev_image, cdna_input, num_masks, color_channels):
  """Apply convolutional dynamic neural advection to previous image.

  Args:
    prev_image: previous image to be transformed.
    cdna_input: hidden lyaer to be used for computing CDNA kernels.
    num_masks: the number of masks and hence the number of CDNA transformations.
    color_channels: the number of color channels in the images.
  Returns:
    List of images transformed by the predicted CDNA kernels.
  """
  batch_size = int(cdna_input.get_shape()[0])
  height = int(prev_image.get_shape()[1])
  width = int(prev_image.get_shape()[2])

  # Predict kernels using linear function of last hidden layer.
  cdna_kerns = slim.layers.fully_connected(
      cdna_input,
      DNA_KERN_SIZE * DNA_KERN_SIZE * num_masks,
      scope='cdna_params',
      activation_fn=None)

  # Reshape and normalize.
  cdna_kerns = tf.reshape(
      cdna_kerns, [batch_size, DNA_KERN_SIZE, DNA_KERN_SIZE, 1, num_masks])
  cdna_kerns = tf.nn.relu(cdna_kerns - RELU_SHIFT) + RELU_SHIFT
  norm_factor = tf.reduce_sum(cdna_kerns, [1, 2, 3], keep_dims=True)
  cdna_kerns /= norm_factor

  # Treat the color channel dimension as the batch dimension since the same
  # transformation is applied to each color channel.
  # Treat the batch dimension as the channel dimension so that
  # depthwise_conv2d can apply a different transformation to each sample.
  cdna_kerns = tf.transpose(cdna_kerns, [1, 2, 0, 4, 3])
  cdna_kerns = tf.reshape(cdna_kerns, [DNA_KERN_SIZE, DNA_KERN_SIZE, batch_size, num_masks])
  # Swap the batch and channel dimensions.
  prev_image = tf.transpose(prev_image, [3, 1, 2, 0])

  # Transform image.
  transformed = tf.nn.depthwise_conv2d(prev_image, cdna_kerns, [1, 1, 1, 1], 'SAME')

  # Transpose the dimensions to where they belong.
  transformed = tf.reshape(transformed, [color_channels, height, width, batch_size, num_masks])
  transformed = tf.transpose(transformed, [3, 1, 2, 0, 4])
  transformed = tf.unstack(transformed, axis=-1)
  return transformed


def dna_transformation(prev_image, dna_input):
  """Apply dynamic neural advection to previous image.

  Args:
    prev_image: previous image to be transformed.
    dna_input: hidden lyaer to be used for computing DNA transformation.
  Returns:
    List of images transformed by the predicted CDNA kernels.
  """
  # Construct translated images.
  prev_image_pad = tf.pad(prev_image, [[0, 0], [2, 2], [2, 2], [0, 0]])
  image_height = int(prev_image.get_shape()[1])
  image_width = int(prev_image.get_shape()[2])

  inputs = []
  for xkern in range(DNA_KERN_SIZE):
    for ykern in range(DNA_KERN_SIZE):
      inputs.append(
          tf.expand_dims(
              tf.slice(prev_image_pad, [0, xkern, ykern, 0],
                       [-1, image_height, image_width, -1]), [3]))
  inputs = tf.concat(axis=3, values=inputs)

  # Normalize channels to 1.
  kernel = tf.nn.relu(dna_input - RELU_SHIFT) + RELU_SHIFT
  kernel = tf.expand_dims(
      kernel / tf.reduce_sum(
          kernel, [3], keep_dims=True), [4])
  return tf.reduce_sum(kernel * inputs, [3], keep_dims=False)


def scheduled_sample(ground_truth_x, generated_x, batch_size, num_ground_truth):
  """Sample batch with specified mix of ground truth and generated data points.

  Args:
    ground_truth_x: tensor of ground-truth data points.
    generated_x: tensor of generated data points.
    batch_size: batch size
    num_ground_truth: number of ground-truth examples to include in batch.
  Returns:
    New batch with num_ground_truth sampled from ground_truth_x and the rest
    from generated_x.
  """
  idx = tf.random_shuffle(tf.range(int(batch_size)))
  ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
  generated_idx = tf.gather(idx, tf.range(num_ground_truth, int(batch_size)))

  ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
  generated_examps = tf.gather(generated_x, generated_idx)
  return tf.dynamic_stitch([ground_truth_idx, generated_idx],
                           [ground_truth_examps, generated_examps])

[((395, 15, 399, 25), 'tensorflow.contrib.slim.layers.fully_connected', 'slim.layers.fully_connected', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((402, 15, 403, 75), 'tensorflow.reshape', 'tf.reshape', ({(403, 6, 403, 16): 'cdna_kerns', (403, 18, 403, 74): '[batch_size, DNA_KERN_SIZE, DNA_KERN_SIZE, 1, num_masks]'}, {}), '(cdna_kerns, [batch_size, DNA_KERN_SIZE, DNA_KERN_SIZE, 1, num_masks]\n    )', True, 'import tensorflow as tf\n'), ((405, 16, 405, 68), 'tensorflow.reduce_sum', 'tf.reduce_sum', (), '', True, 'import tensorflow as tf\n'), ((412, 15, 412, 56), 'tensorflow.transpose', 'tf.transpose', ({(412, 28, 412, 38): 'cdna_kerns', (412, 40, 412, 55): '[1, 2, 0, 4, 3]'}, {}), '(cdna_kerns, [1, 2, 0, 4, 3])', True, 'import tensorflow as tf\n'), ((413, 15, 413, 92), 'tensorflow.reshape', 'tf.reshape', ({(413, 26, 413, 36): 'cdna_kerns', (413, 38, 413, 91): '[DNA_KERN_SIZE, DNA_KERN_SIZE, batch_size, num_masks]'}, {}), '(cdna_kerns, [DNA_KERN_SIZE, DNA_KERN_SIZE, batch_size, num_masks])', True, 'import tensorflow as tf\n'), ((415, 15, 415, 53), 'tensorflow.transpose', 'tf.transpose', ({(415, 28, 415, 38): 'prev_image', (415, 40, 415, 52): '[3, 1, 2, 0]'}, {}), '(prev_image, [3, 1, 2, 0])', True, 'import tensorflow as tf\n'), ((418, 16, 418, 84), 'tensorflow.nn.depthwise_conv2d', 'tf.nn.depthwise_conv2d', ({(418, 39, 418, 49): 'prev_image', (418, 51, 418, 61): 'cdna_kerns', (418, 63, 418, 75): '[1, 1, 1, 1]', (418, 77, 418, 83): '"""SAME"""'}, {}), "(prev_image, cdna_kerns, [1, 1, 1, 1], 'SAME')", True, 'import tensorflow as tf\n'), ((421, 16, 421, 95), 'tensorflow.reshape', 'tf.reshape', ({(421, 27, 421, 38): 'transformed', (421, 40, 421, 94): '[color_channels, height, width, batch_size, num_masks]'}, {}), '(transformed, [color_channels, height, width, batch_size, num_masks])', True, 'import tensorflow as tf\n'), ((422, 16, 422, 58), 'tensorflow.transpose', 'tf.transpose', ({(422, 29, 422, 40): 'transformed', (422, 42, 422, 57): '[3, 1, 2, 0, 4]'}, {}), '(transformed, [3, 1, 2, 0, 4])', True, 'import tensorflow as tf\n'), ((423, 16, 423, 48), 'tensorflow.unstack', 'tf.unstack', (), '', True, 'import tensorflow as tf\n'), ((437, 19, 437, 71), 'tensorflow.pad', 'tf.pad', ({(437, 26, 437, 36): 'prev_image', (437, 38, 437, 70): '[[0, 0], [2, 2], [2, 2], [0, 0]]'}, {}), '(prev_image, [[0, 0], [2, 2], [2, 2], [0, 0]])', True, 'import tensorflow as tf\n'), ((448, 11, 448, 43), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((455, 9, 455, 61), 'tensorflow.reduce_sum', 'tf.reduce_sum', (), '', True, 'import tensorflow as tf\n'), ((474, 24, 474, 67), 'tensorflow.gather', 'tf.gather', ({(474, 34, 474, 48): 'ground_truth_x', (474, 50, 474, 66): 'ground_truth_idx'}, {}), '(ground_truth_x, ground_truth_idx)', True, 'import tensorflow as tf\n'), ((475, 21, 475, 58), 'tensorflow.gather', 'tf.gather', ({(475, 31, 475, 42): 'generated_x', (475, 44, 475, 57): 'generated_idx'}, {}), '(generated_x, generated_idx)', True, 'import tensorflow as tf\n'), ((476, 9, 477, 67), 'tensorflow.dynamic_stitch', 'tf.dynamic_stitch', ({(476, 27, 476, 60): '[ground_truth_idx, generated_idx]', (477, 27, 477, 66): '[ground_truth_examps, generated_examps]'}, {}), '([ground_truth_idx, generated_idx], [ground_truth_examps,\n    generated_examps])', True, 'import tensorflow as tf\n'), ((69, 7, 69, 49), 'tensorflow.contrib.slim.arg_scope', 'slim.arg_scope', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((70, 21, 70, 41), 'tensorflow.concat', 'tf.concat', ({(70, 31, 70, 37): 'images', (70, 39, 70, 40): '3'}, {}), '(images, 3)', True, 'import tensorflow as tf\n'), ((72, 18, 78, 52), 'tensorflow.contrib.slim.conv2d', 'slim.conv2d', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((80, 18, 86, 52), 'tensorflow.contrib.slim.conv2d', 'slim.conv2d', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((88, 18, 94, 52), 'tensorflow.contrib.slim.conv2d', 'slim.conv2d', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((96, 18, 103, 56), 'tensorflow.contrib.slim.conv2d', 'slim.conv2d', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((105, 17, 111, 55), 'tensorflow.contrib.slim.conv2d', 'slim.conv2d', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((116, 18, 116, 44), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(116, 33, 116, 43): 'divergence'}, {}), '(divergence)', True, 'import tensorflow as tf\n'), ((120, 14, 122, 25), 'tensorflow.random_normal', 'tf.random_normal', (), '', True, 'import tensorflow as tf\n'), ((125, 14, 125, 73), 'tensorflow.random_normal', 'tf.random_normal', (), '', True, 'import tensorflow as tf\n'), ((172, 12, 172, 30), 'tensorflow.identity', 'tf.identity', ({(172, 24, 172, 29): 'image'}, {}), '(image)', True, 'import tensorflow as tf\n'), ((193, 23, 193, 62), 'numpy.array', 'np.array', ({(193, 32, 193, 61): '[32, 32, 64, 64, 128, 64, 32]'}, {}), '([32, 32, 64, 64, 128, 64, 32])', True, 'import numpy as np\n'), ((368, 6, 368, 58), 'numpy.array', 'np.array', ({(368, 15, 368, 45): '[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]', (368, 47, 368, 57): 'np.float32'}, {}), '([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], np.float32)', True, 'import numpy as np\n'), ((404, 15, 404, 50), 'tensorflow.nn.relu', 'tf.nn.relu', ({(404, 26, 404, 49): '(cdna_kerns - RELU_SHIFT)'}, {}), '(cdna_kerns - RELU_SHIFT)', True, 'import tensorflow as tf\n'), ((451, 11, 451, 45), 'tensorflow.nn.relu', 'tf.nn.relu', ({(451, 22, 451, 44): '(dna_input - RELU_SHIFT)'}, {}), '(dna_input - RELU_SHIFT)', True, 'import tensorflow as tf\n'), ((471, 36, 471, 62), 'tensorflow.range', 'tf.range', ({(471, 45, 471, 61): 'num_ground_truth'}, {}), '(num_ground_truth)', True, 'import tensorflow as tf\n'), ((209, 9, 212, 20), 'tensorflow.contrib.slim.arg_scope', 'slim.arg_scope', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((226, 21, 226, 70), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((228, 13, 234, 53), 'tensorflow.contrib.slim.layers.conv2d', 'slim.layers.conv2d', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((238, 16, 238, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((241, 16, 241, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((247, 16, 247, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((250, 16, 250, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((280, 16, 280, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((286, 16, 286, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((288, 16, 288, 57), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((294, 16, 294, 66), 'tensorflow.contrib.layers.python.layers.layer_norm', 'tf_layers.layer_norm', (), '', True, 'from tensorflow.contrib.layers.python import layers as tf_layers\n'), ((297, 16, 297, 57), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((332, 14, 333, 79), 'tensorflow.contrib.slim.layers.conv2d_transpose', 'slim.layers.conv2d_transpose', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((337, 18, 337, 81), 'tensorflow.split', 'tf.split', (), '', True, 'import tensorflow as tf\n'), ((374, 23, 374, 54), 'spatial_transformer.transformer', 'transformer', ({(374, 35, 374, 45): 'prev_image', (374, 47, 374, 53): 'params'}, {}), '(prev_image, params)', False, 'from spatial_transformer import transformer\n'), ((453, 15, 454, 38), 'tensorflow.reduce_sum', 'tf.reduce_sum', (), '', True, 'import tensorflow as tf\n'), ((45, 62, 45, 79), 'tensorflow.exp', 'tf.exp', ({(45, 69, 45, 78): 'log_sigma'}, {}), '(log_sigma)', True, 'import tensorflow as tf\n'), ((261, 15, 261, 54), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((307, 15, 308, 84), 'tensorflow.contrib.slim.layers.conv2d_transpose', 'slim.layers.conv2d_transpose', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((311, 15, 312, 82), 'tensorflow.contrib.slim.layers.conv2d_transpose', 'slim.layers.conv2d_transpose', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((319, 21, 320, 44), 'tensorflow.contrib.slim.layers.fully_connected', 'slim.layers.fully_connected', (), '', True, 'import tensorflow.contrib.slim as slim\n'), ((45, 46, 45, 59), 'tensorflow.square', 'tf.square', ({(45, 56, 45, 58): 'mu'}, {}), '(mu)', True, 'import tensorflow as tf\n'), ((189, 17, 189, 40), 'tensorflow.to_float', 'tf.to_float', ({(189, 29, 189, 39): 'batch_size'}, {}), '(batch_size)', True, 'import tensorflow as tf\n'), ((272, 13, 272, 46), 'tensorflow.control_dependencies', 'tf.control_dependencies', ({(272, 37, 272, 45): '[latent]'}, {}), '([latent])', True, 'import tensorflow as tf\n'), ((273, 17, 273, 45), 'tensorflow.concat', 'tf.concat', ({(273, 27, 273, 41): '[enc2, latent]', (273, 43, 273, 44): '3'}, {}), '([enc2, latent], 3)', True, 'import tensorflow as tf\n'), ((315, 23, 315, 42), 'tensorflow.nn.sigmoid', 'tf.nn.sigmoid', ({(315, 37, 315, 41): 'enc7'}, {}), '(enc7)', True, 'import tensorflow as tf\n'), ((335, 24, 335, 62), 'tensorflow.reshape', 'tf.reshape', ({(335, 35, 335, 40): 'masks', (335, 42, 335, 61): '[-1, num_masks + 1]'}, {}), '(masks, [-1, num_masks + 1])', True, 'import tensorflow as tf\n'), ((446, 14, 447, 59), 'tensorflow.slice', 'tf.slice', ({(446, 23, 446, 37): 'prev_image_pad', (446, 39, 446, 59): '[0, xkern, ykern, 0]', (447, 23, 447, 58): '[-1, image_height, image_width, -1]'}, {}), '(prev_image_pad, [0, xkern, ykern, 0], [-1, image_height,\n    image_width, -1])', True, 'import tensorflow as tf\n'), ((189, 53, 189, 73), 'tensorflow.exp', 'tf.exp', ({(189, 60, 189, 72): 'iter_num / k'}, {}), '(iter_num / k)', True, 'import tensorflow as tf\n'), ((270, 35, 270, 54), 'tensorflow.identity', 'tf.identity', ({(270, 47, 270, 53): 'latent'}, {}), '(latent)', True, 'import tensorflow as tf\n'), ((271, 49, 271, 73), 'tensorflow.exp', 'tf.exp', ({(271, 56, 271, 72): 'latent_std / 2.0'}, {}), '(latent_std / 2.0)', True, 'import tensorflow as tf\n')]