text: # Copyright 2015 Google Inc. 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.
  # ==============================================================================
  """The Neural GPU Model."""
  
  import time
  
  import numpy as np
  from six.moves import xrange
  import tensorflow as tf
  
  from tensorflow.python.framework import function
  import data_utils as data
  
  do_jit = False  # Gives more speed but experimental for now.
  jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
  
  
  def conv_linear(args, kw, kh, nin, nout, rate, do_bias, bias_start, prefix):
    """Convolutional linear map."""
    if not isinstance(args, (list, tuple)):
      args = [args]
    with tf.variable_scope(prefix):
      with tf.device("/cpu:0"):
        k = tf.get_variable("CvK", [kw, kh, nin, nout])
      if len(args) == 1:
        arg = args[0]
      else:
        arg = tf.concat(axis=3, values=args)
      res = tf.nn.convolution(arg, k, dilation_rate=(rate, 1), padding="SAME")
      if not do_bias: return res
      with tf.device("/cpu:0"):
        bias_term = tf.get_variable(
            "CvB", [nout], initializer=tf.constant_initializer(bias_start))
      bias_term = tf.reshape(bias_term, [1, 1, 1, nout])
      return res + bias_term
  
  
  def sigmoid_cutoff(x, cutoff):
    """Sigmoid with cutoff, e.g., 1.2sigmoid(x) - 0.1."""
    y = tf.sigmoid(x)
    if cutoff < 1.01: return y
    d = (cutoff - 1.0) / 2.0
    return tf.minimum(1.0, tf.maximum(0.0, cutoff * y - d), name="cutoff_min")
  
  
  @function.Defun(tf.float32, noinline=True)
  def sigmoid_cutoff_12(x):
    """Sigmoid with cutoff 1.2, specialized for speed and memory use."""
    y = tf.sigmoid(x)
    return tf.minimum(1.0, tf.maximum(0.0, 1.2 * y - 0.1), name="cutoff_min_12")
  
  
  @function.Defun(tf.float32, noinline=True)
  def sigmoid_hard(x):
    """Hard sigmoid."""
    return tf.minimum(1.0, tf.maximum(0.0, 0.25 * x + 0.5))
  
  
  def place_at14(decided, selected, it):
    """Place selected at it-th coordinate of decided, dim=1 of 4."""
    slice1 = decided[:, :it, :, :]
    slice2 = decided[:, it + 1:, :, :]
    return tf.concat(axis=1, values=[slice1, selected, slice2])
  
  
  def place_at13(decided, selected, it):
    """Place selected at it-th coordinate of decided, dim=1 of 3."""
    slice1 = decided[:, :it, :]
    slice2 = decided[:, it + 1:, :]
    return tf.concat(axis=1, values=[slice1, selected, slice2])
  
  
  def tanh_cutoff(x, cutoff):
    """Tanh with cutoff, e.g., 1.1tanh(x) cut to [-1. 1]."""
    y = tf.tanh(x)
    if cutoff < 1.01: return y
    d = (cutoff - 1.0) / 2.0
    return tf.minimum(1.0, tf.maximum(-1.0, (1.0 + d) * y))
  
  
  @function.Defun(tf.float32, noinline=True)
  def tanh_hard(x):
    """Hard tanh."""
    return tf.minimum(1.0, tf.maximum(0.0, x))
  
  
  def layer_norm(x, nmaps, prefix, epsilon=1e-5):
    """Layer normalize the 4D tensor x, averaging over the last dimension."""
    with tf.variable_scope(prefix):
      scale = tf.get_variable("layer_norm_scale", [nmaps],
                              initializer=tf.ones_initializer())
      bias = tf.get_variable("layer_norm_bias", [nmaps],
                             initializer=tf.zeros_initializer())
      mean, variance = tf.nn.moments(x, [3], keep_dims=True)
      norm_x = (x - mean) / tf.sqrt(variance + epsilon)
      return norm_x * scale + bias
  
  
  def conv_gru(inpts, mem, kw, kh, nmaps, rate, cutoff, prefix, do_layer_norm,
               args_len=None):
    """Convolutional GRU."""
    def conv_lin(args, suffix, bias_start):
      total_args_len = args_len or len(args) * nmaps
      res = conv_linear(args, kw, kh, total_args_len, nmaps, rate, True,
                        bias_start, prefix + "/" + suffix)
      if do_layer_norm:
        return layer_norm(res, nmaps, prefix + "/" + suffix)
      else:
        return res
    if cutoff == 1.2:
      reset = sigmoid_cutoff_12(conv_lin(inpts + [mem], "r", 1.0))
      gate = sigmoid_cutoff_12(conv_lin(inpts + [mem], "g", 1.0))
    elif cutoff > 10:
      reset = sigmoid_hard(conv_lin(inpts + [mem], "r", 1.0))
      gate = sigmoid_hard(conv_lin(inpts + [mem], "g", 1.0))
    else:
      reset = sigmoid_cutoff(conv_lin(inpts + [mem], "r", 1.0), cutoff)
      gate = sigmoid_cutoff(conv_lin(inpts + [mem], "g", 1.0), cutoff)
    if cutoff > 10:
      candidate = tanh_hard(conv_lin(inpts + [reset * mem], "c", 0.0))
    else:
      # candidate = tanh_cutoff(conv_lin(inpts + [reset * mem], "c", 0.0), cutoff)
      candidate = tf.tanh(conv_lin(inpts + [reset * mem], "c", 0.0))
    return gate * mem + (1 - gate) * candidate
  
  
  CHOOSE_K = 256
  
  
  def memory_call(q, l, nmaps, mem_size, vocab_size, num_gpus, update_mem):
    raise ValueError("Fill for experiments with additional memory structures.")
  
  
  def memory_run(step, nmaps, mem_size, batch_size, vocab_size,
                 global_step, do_training, update_mem, decay_factor, num_gpus,
                 target_emb_weights, output_w, gpu_targets_tn, it):
    """Run memory."""
    q = step[:, 0, it, :]
    mlabels = gpu_targets_tn[:, it, 0]
    res, mask, mem_loss = memory_call(
        q, mlabels, nmaps, mem_size, vocab_size, num_gpus, update_mem)
    res = tf.gather(target_emb_weights, res) * tf.expand_dims(mask[:, 0], 1)
  
    # Mix gold and original in the first steps, 20% later.
    gold = tf.nn.dropout(tf.gather(target_emb_weights, mlabels), 0.7)
    use_gold = 1.0 - tf.cast(global_step, tf.float32) / (1000. * decay_factor)
    use_gold = tf.maximum(use_gold, 0.2) * do_training
    mem = tf.cond(tf.less(tf.random_uniform([]), use_gold),
                  lambda: use_gold * gold + (1.0 - use_gold) * res,
                  lambda: res)
    mem = tf.reshape(mem, [-1, 1, 1, nmaps])
    return mem, mem_loss, update_mem
  
  
  @tf.RegisterGradient("CustomIdG")
  def _custom_id_grad(_, grads):
    return grads
  
  
  def quantize(t, quant_scale, max_value=1.0):
    """Quantize a tensor t with each element in [-max_value, max_value]."""
    t = tf.minimum(max_value, tf.maximum(t, -max_value))
    big = quant_scale * (t + max_value) + 0.5
    with tf.get_default_graph().gradient_override_map({"Floor": "CustomIdG"}):
      res = (tf.floor(big) / quant_scale) - max_value
    return res
  
  
  def quantize_weights_op(quant_scale, max_value):
    ops = [v.assign(quantize(v, quant_scale, float(max_value)))
           for v in tf.trainable_variables()]
    return tf.group(*ops)
  
  
  def autoenc_quantize(x, nbits, nmaps, do_training, layers=1):
    """Autoencoder into nbits vectors of bits, using noise and sigmoids."""
    enc_x = tf.reshape(x, [-1, nmaps])
    for i in xrange(layers - 1):
      enc_x = tf.layers.dense(enc_x, nmaps, name="autoenc_%d" % i)
    enc_x = tf.layers.dense(enc_x, nbits, name="autoenc_%d" % (layers - 1))
    noise = tf.truncated_normal(tf.shape(enc_x), stddev=2.0)
    dec_x = sigmoid_cutoff_12(enc_x + noise * do_training)
    dec_x = tf.reshape(dec_x, [-1, nbits])
    for i in xrange(layers):
      dec_x = tf.layers.dense(dec_x, nmaps, name="autodec_%d" % i)
    return tf.reshape(dec_x, tf.shape(x))
  
  
  def make_dense(targets, noclass, low_param):
    """Move a batch of targets to a dense 1-hot representation."""
    low = low_param / float(noclass - 1)
    high = 1.0 - low * (noclass - 1)
    targets = tf.cast(targets, tf.int64)
    return tf.one_hot(targets, depth=noclass, on_value=high, off_value=low)
  
  
  def reorder_beam(beam_size, batch_size, beam_val, output, is_first,
                   tensors_to_reorder):
    """Reorder to minimize beam costs."""
    # beam_val is [batch_size x beam_size]; let b = batch_size * beam_size
    # decided is len x b x a x b
    # output is b x out_size; step is b x len x a x b;
    outputs = tf.split(axis=0, num_or_size_splits=beam_size, value=tf.nn.log_softmax(output))
    all_beam_vals, all_beam_idx = [], []
    beam_range = 1 if is_first else beam_size
    for i in xrange(beam_range):
      top_out, top_out_idx = tf.nn.top_k(outputs[i], k=beam_size)
      cur_beam_val = beam_val[:, i]
      top_out = tf.Print(top_out, [top_out, top_out_idx, beam_val, i,
                                   cur_beam_val], "GREPO", summarize=8)
      all_beam_vals.append(top_out + tf.expand_dims(cur_beam_val, 1))
      all_beam_idx.append(top_out_idx)
    all_beam_idx = tf.reshape(tf.transpose(tf.concat(axis=1, values=all_beam_idx), [1, 0]),
                              [-1])
    top_beam, top_beam_idx = tf.nn.top_k(tf.concat(axis=1, values=all_beam_vals), k=beam_size)
    top_beam_idx = tf.Print(top_beam_idx, [top_beam, top_beam_idx],
                            "GREP", summarize=8)
    reordered = [[] for _ in xrange(len(tensors_to_reorder) + 1)]
    top_out_idx = []
    for i in xrange(beam_size):
      which_idx = top_beam_idx[:, i] * batch_size + tf.range(batch_size)
      top_out_idx.append(tf.gather(all_beam_idx, which_idx))
      which_beam = top_beam_idx[:, i] / beam_size  # [batch]
      which_beam = which_beam * batch_size + tf.range(batch_size)
      reordered[0].append(tf.gather(output, which_beam))
      for i, t in enumerate(tensors_to_reorder):
        reordered[i + 1].append(tf.gather(t, which_beam))
    new_tensors = [tf.concat(axis=0, values=t) for t in reordered]
    top_out_idx = tf.concat(axis=0, values=top_out_idx)
    return (top_beam, new_tensors[0], top_out_idx, new_tensors[1:])
  
  
  class NeuralGPU(object):
    """Neural GPU Model."""
  
    def __init__(self, nmaps, vec_size, niclass, noclass, dropout,
                 max_grad_norm, cutoff, nconvs, kw, kh, height, mem_size,
                 learning_rate, min_length, num_gpus, num_replicas,
                 grad_noise_scale, sampling_rate, act_noise=0.0, do_rnn=False,
                 atrous=False, beam_size=1, backward=True, do_layer_norm=False,
                 autoenc_decay=1.0):
      # Feeds for parameters and ops to update them.
      self.nmaps = nmaps
      if backward:
        self.global_step = tf.Variable(0, trainable=False, name="global_step")
        self.cur_length = tf.Variable(min_length, trainable=False)
        self.cur_length_incr_op = self.cur_length.assign_add(1)
        self.lr = tf.Variable(learning_rate, trainable=False)
        self.lr_decay_op = self.lr.assign(self.lr * 0.995)
      self.do_training = tf.placeholder(tf.float32, name="do_training")
      self.update_mem = tf.placeholder(tf.int32, name="update_mem")
      self.noise_param = tf.placeholder(tf.float32, name="noise_param")
  
      # Feeds for inputs, targets, outputs, losses, etc.
      self.input = tf.placeholder(tf.int32, name="inp")
      self.target = tf.placeholder(tf.int32, name="tgt")
      self.prev_step = tf.placeholder(tf.float32, name="prev_step")
      gpu_input = tf.split(axis=0, num_or_size_splits=num_gpus, value=self.input)
      gpu_target = tf.split(axis=0, num_or_size_splits=num_gpus, value=self.target)
      gpu_prev_step = tf.split(axis=0, num_or_size_splits=num_gpus, value=self.prev_step)
      batch_size = tf.shape(gpu_input[0])[0]
  
      if backward:
        adam_lr = 0.005 * self.lr
        adam = tf.train.AdamOptimizer(adam_lr, epsilon=1e-3)
  
        def adam_update(grads):
          return adam.apply_gradients(zip(grads, tf.trainable_variables()),
                                      global_step=self.global_step,
                                      name="adam_update")
  
      # When switching from Adam to SGD we perform reverse-decay.
      if backward:
        global_step_float = tf.cast(self.global_step, tf.float32)
        sampling_decay_exponent = global_step_float / 100000.0
        sampling_decay = tf.maximum(0.05, tf.pow(0.5, sampling_decay_exponent))
        self.sampling = sampling_rate * 0.05 / sampling_decay
      else:
        self.sampling = tf.constant(0.0)
  
      # Cache variables on cpu if needed.
      if num_replicas > 1 or num_gpus > 1:
        with tf.device("/cpu:0"):
          caching_const = tf.constant(0)
        tf.get_variable_scope().set_caching_device(caching_const.op.device)
        # partitioner = tf.variable_axis_size_partitioner(1024*256*4)
        # tf.get_variable_scope().set_partitioner(partitioner)
  
      def gpu_avg(l):
        if l[0] is None:
          for elem in l:
            assert elem is None
          return 0.0
        if len(l) < 2:
          return l[0]
        return sum(l) / float(num_gpus)
  
      self.length_tensor = tf.placeholder(tf.int32, name="length")
  
      with tf.device("/cpu:0"):
        emb_weights = tf.get_variable(
            "embedding", [niclass, vec_size],
            initializer=tf.random_uniform_initializer(-1.7, 1.7))
        if beam_size > 0:
          target_emb_weights = tf.get_variable(
              "target_embedding", [noclass, nmaps],
              initializer=tf.random_uniform_initializer(-1.7, 1.7))
        e0 = tf.scatter_update(emb_weights,
                               tf.constant(0, dtype=tf.int32, shape=[1]),
                               tf.zeros([1, vec_size]))
        output_w = tf.get_variable("output_w", [nmaps, noclass], tf.float32)
  
      def conv_rate(layer):
        if atrous:
          return 2**layer
        return 1
  
      # pylint: disable=cell-var-from-loop
      def enc_step(step):
        """Encoder step."""
        if autoenc_decay < 1.0:
          quant_step = autoenc_quantize(step, 16, nmaps, self.do_training)
          if backward:
            exp_glob = tf.train.exponential_decay(1.0, self.global_step - 10000,
                                                  1000, autoenc_decay)
            dec_factor = 1.0 - exp_glob  # * self.do_training
            dec_factor = tf.cond(tf.less(self.global_step, 10500),
                                 lambda: tf.constant(0.05), lambda: dec_factor)
          else:
            dec_factor = 1.0
          cur = tf.cond(tf.less(tf.random_uniform([]), dec_factor),
                        lambda: quant_step, lambda: step)
        else:
          cur = step
        if dropout > 0.0001:
          cur = tf.nn.dropout(cur, keep_prob)
        if act_noise > 0.00001:
          cur += tf.truncated_normal(tf.shape(cur)) * act_noise_scale
        # Do nconvs-many CGRU steps.
        if do_jit and tf.get_variable_scope().reuse:
          with jit_scope():
            for layer in xrange(nconvs):
              cur = conv_gru([], cur, kw, kh, nmaps, conv_rate(layer),
                             cutoff, "ecgru_%d" % layer, do_layer_norm)
        else:
          for layer in xrange(nconvs):
            cur = conv_gru([], cur, kw, kh, nmaps, conv_rate(layer),
                           cutoff, "ecgru_%d" % layer, do_layer_norm)
        return cur
  
      zero_tgt = tf.zeros([batch_size, nmaps, 1])
      zero_tgt.set_shape([None, nmaps, 1])
  
      def dec_substep(step, decided):
        """Decoder sub-step."""
        cur = step
        if dropout > 0.0001:
          cur = tf.nn.dropout(cur, keep_prob)
        if act_noise > 0.00001:
          cur += tf.truncated_normal(tf.shape(cur)) * act_noise_scale
        # Do nconvs-many CGRU steps.
        if do_jit and tf.get_variable_scope().reuse:
          with jit_scope():
            for layer in xrange(nconvs):
              cur = conv_gru([decided], cur, kw, kh, nmaps, conv_rate(layer),
                             cutoff, "dcgru_%d" % layer, do_layer_norm)
        else:
          for layer in xrange(nconvs):
            cur = conv_gru([decided], cur, kw, kh, nmaps, conv_rate(layer),
                           cutoff, "dcgru_%d" % layer, do_layer_norm)
        return cur
      # pylint: enable=cell-var-from-loop
  
      def dec_step(step, it, it_int, decided, output_ta, tgts,
                   mloss, nupd_in, out_idx, beam_cost):
        """Decoder step."""
        nupd, mem_loss = 0, 0.0
        if mem_size > 0:
          it_incr = tf.minimum(it+1, length - 1)
          mem, mem_loss, nupd = memory_run(
              step, nmaps, mem_size, batch_size, noclass, self.global_step,
              self.do_training, self.update_mem, 10, num_gpus,
              target_emb_weights, output_w, gpu_targets_tn, it_incr)
        step = dec_substep(step, decided)
        output_l = tf.expand_dims(tf.expand_dims(step[:, it, 0, :], 1), 1)
        # Calculate argmax output.
        output = tf.reshape(output_l, [-1, nmaps])
        # pylint: disable=cell-var-from-loop
        output = tf.matmul(output, output_w)
        if beam_size > 1:
          beam_cost, output, out, reordered = reorder_beam(
              beam_size, batch_size, beam_cost, output, it_int == 0,
              [output_l, out_idx, step, decided])
          [output_l, out_idx, step, decided] = reordered
        else:
          # Scheduled sampling.
          out = tf.multinomial(tf.stop_gradient(output), 1)
          out = tf.to_int32(tf.squeeze(out, [1]))
        out_write = output_ta.write(it, output_l[:batch_size, :, :, :])
        output = tf.gather(target_emb_weights, out)
        output = tf.reshape(output, [-1, 1, nmaps])
        output = tf.concat(axis=1, values=[output] * height)
        tgt = tgts[it, :, :, :]
        selected = tf.cond(tf.less(tf.random_uniform([]), self.sampling),
                           lambda: output, lambda: tgt)
        # pylint: enable=cell-var-from-loop
        dec_write = place_at14(decided, tf.expand_dims(selected, 1), it)
        out_idx = place_at13(
            out_idx, tf.reshape(out, [beam_size * batch_size, 1, 1]), it)
        if mem_size > 0:
          mem = tf.concat(axis=2, values=[mem] * height)
          dec_write = place_at14(dec_write, mem, it_incr)
        return (step, dec_write, out_write, mloss + mem_loss, nupd_in + nupd,
                out_idx, beam_cost)
  
      # Main model construction.
      gpu_outputs = []
      gpu_losses = []
      gpu_grad_norms = []
      grads_list = []
      gpu_out_idx = []
      self.after_enc_step = []
      for gpu in xrange(num_gpus):  # Multi-GPU towers, average gradients later.
        length = self.length_tensor
        length_float = tf.cast(length, tf.float32)
        if gpu > 0:
          tf.get_variable_scope().reuse_variables()
        gpu_outputs.append([])
        gpu_losses.append([])
        gpu_grad_norms.append([])
        with tf.name_scope("gpu%d" % gpu), tf.device("/gpu:%d" % gpu):
          # Main graph creation loop.
          data.print_out("Creating model.")
          start_time = time.time()
  
          # Embed inputs and calculate mask.
          with tf.device("/cpu:0"):
            tgt_shape = tf.shape(tf.squeeze(gpu_target[gpu], [1]))
            weights = tf.where(tf.squeeze(gpu_target[gpu], [1]) > 0,
                               tf.ones(tgt_shape), tf.zeros(tgt_shape))
  
            # Embed inputs and targets.
            with tf.control_dependencies([e0]):
              start = tf.gather(emb_weights, gpu_input[gpu])  # b x h x l x nmaps
              gpu_targets_tn = gpu_target[gpu]  # b x 1 x len
              if beam_size > 0:
                embedded_targets_tn = tf.gather(target_emb_weights,
                                                gpu_targets_tn)
                embedded_targets_tn = tf.transpose(
                    embedded_targets_tn, [2, 0, 1, 3])  # len x b x 1 x nmaps
                embedded_targets_tn = tf.concat(axis=2, values=[embedded_targets_tn] * height)
  
          # First image comes from start by applying convolution and adding 0s.
          start = tf.transpose(start, [0, 2, 1, 3])  # Now b x len x h x vec_s
          first = conv_linear(start, 1, 1, vec_size, nmaps, 1, True, 0.0, "input")
          first = layer_norm(first, nmaps, "input")
  
          # Computation steps.
          keep_prob = dropout * 3.0 / tf.sqrt(length_float)
          keep_prob = 1.0 - self.do_training * keep_prob
          act_noise_scale = act_noise * self.do_training
  
          # Start with a convolutional gate merging previous step.
          step = conv_gru([gpu_prev_step[gpu]], first,
                          kw, kh, nmaps, 1, cutoff, "first", do_layer_norm)
  
          # This is just for running a baseline RNN seq2seq model.
          if do_rnn:
            self.after_enc_step.append(step)  # Not meaningful here, but needed.
            def lstm_cell():
              return tf.contrib.rnn.BasicLSTMCell(height * nmaps)
            cell = tf.contrib.rnn.MultiRNNCell(
                [lstm_cell() for _ in range(nconvs)])
            with tf.variable_scope("encoder"):
              encoder_outputs, encoder_state = tf.nn.dynamic_rnn(
                  cell, tf.reshape(step, [batch_size, length, height * nmaps]),
                  dtype=tf.float32, time_major=False)
  
            # Attention.
            attn = tf.layers.dense(
                encoder_outputs, height * nmaps, name="attn1")
  
            # pylint: disable=cell-var-from-loop
            @function.Defun(noinline=True)
            def attention_query(query, attn_v):
              vecs = tf.tanh(attn + tf.expand_dims(query, 1))
              mask = tf.reduce_sum(vecs * tf.reshape(attn_v, [1, 1, -1]), 2)
              mask = tf.nn.softmax(mask)
              return tf.reduce_sum(encoder_outputs * tf.expand_dims(mask, 2), 1)
  
            with tf.variable_scope("decoder"):
              def decoder_loop_fn(state__prev_cell_out__unused, cell_inp__cur_tgt):
                """Decoder loop function."""
                state, prev_cell_out, _ = state__prev_cell_out__unused
                cell_inp, cur_tgt = cell_inp__cur_tgt
                attn_q = tf.layers.dense(prev_cell_out, height * nmaps,
                                         name="attn_query")
                attn_res = attention_query(attn_q, tf.get_variable(
                    "attn_v", [height * nmaps],
                    initializer=tf.random_uniform_initializer(-0.1, 0.1)))
                concatenated = tf.reshape(tf.concat(axis=1, values=[cell_inp, attn_res]),
                                          [batch_size, 2 * height * nmaps])
                cell_inp = tf.layers.dense(
                    concatenated, height * nmaps, name="attn_merge")
                output, new_state = cell(cell_inp, state)
  
                mem_loss = 0.0
                if mem_size > 0:
                  res, mask, mem_loss = memory_call(
                      output, cur_tgt, height * nmaps, mem_size, noclass,
                      num_gpus, self.update_mem)
                  res = tf.gather(target_emb_weights, res)
                  res *= tf.expand_dims(mask[:, 0], 1)
                  output = tf.layers.dense(
                      tf.concat(axis=1, values=[output, res]), height * nmaps, name="rnnmem")
  
                return new_state, output, mem_loss
              # pylint: enable=cell-var-from-loop
              gpu_targets = tf.squeeze(gpu_target[gpu], [1])  # b x len
              gpu_tgt_trans = tf.transpose(gpu_targets, [1, 0])
              dec_zero = tf.zeros([batch_size, 1], dtype=tf.int32)
              dec_inp = tf.concat(axis=1, values=[dec_zero, gpu_targets])
              dec_inp = dec_inp[:, :length]
              embedded_dec_inp = tf.gather(target_emb_weights, dec_inp)
              embedded_dec_inp_proj = tf.layers.dense(
                  embedded_dec_inp, height * nmaps, name="dec_proj")
              embedded_dec_inp_proj = tf.transpose(embedded_dec_inp_proj,
                                                   [1, 0, 2])
              init_vals = (encoder_state,
                           tf.zeros([batch_size, height * nmaps]), 0.0)
              _, dec_outputs, mem_losses = tf.scan(
                  decoder_loop_fn, (embedded_dec_inp_proj, gpu_tgt_trans),
                  initializer=init_vals)
            mem_loss = tf.reduce_mean(mem_losses)
            outputs = tf.layers.dense(dec_outputs, nmaps, name="out_proj")
            # Final convolution to get logits, list outputs.
            outputs = tf.matmul(tf.reshape(outputs, [-1, nmaps]), output_w)
            outputs = tf.reshape(outputs, [length, batch_size, noclass])
            gpu_out_idx.append(tf.argmax(outputs, 2))
          else:  # Here we go with the Neural GPU.
            # Encoder.
            enc_length = length
            step = enc_step(step)  # First step hard-coded.
            # pylint: disable=cell-var-from-loop
            i = tf.constant(1)
            c = lambda i, _s: tf.less(i, enc_length)
            def enc_step_lambda(i, step):
              with tf.variable_scope(tf.get_variable_scope(), reuse=True):
                new_step = enc_step(step)
              return (i + 1, new_step)
            _, step = tf.while_loop(
                c, enc_step_lambda, [i, step],
                parallel_iterations=1, swap_memory=True)
            # pylint: enable=cell-var-from-loop
  
            self.after_enc_step.append(step)
  
            # Decoder.
            if beam_size > 0:
              output_ta = tf.TensorArray(
                  dtype=tf.float32, size=length, dynamic_size=False,
                  infer_shape=False, name="outputs")
              out_idx = tf.zeros([beam_size * batch_size, length, 1],
                                 dtype=tf.int32)
              decided_t = tf.zeros([beam_size * batch_size, length,
                                    height, vec_size])
  
              # Prepare for beam search.
              tgts = tf.concat(axis=1, values=[embedded_targets_tn] * beam_size)
              beam_cost = tf.zeros([batch_size, beam_size])
              step = tf.concat(axis=0, values=[step] * beam_size)
              # First step hard-coded.
              step, decided_t, output_ta, mem_loss, nupd, oi, bc = dec_step(
                  step, 0, 0, decided_t, output_ta, tgts, 0.0, 0, out_idx,
                  beam_cost)
              tf.get_variable_scope().reuse_variables()
              # pylint: disable=cell-var-from-loop
              def step_lambda(i, step, dec_t, out_ta, ml, nu, oi, bc):
                with tf.variable_scope(tf.get_variable_scope(), reuse=True):
                  s, d, t, nml, nu, oi, bc = dec_step(
                      step, i, 1, dec_t, out_ta, tgts, ml, nu, oi, bc)
                return (i + 1, s, d, t, nml, nu, oi, bc)
              i = tf.constant(1)
              c = lambda i, _s, _d, _o, _ml, _nu, _oi, _bc: tf.less(i, length)
              _, step, _, output_ta, mem_loss, nupd, out_idx, _ = tf.while_loop(
                  c, step_lambda,
                  [i, step, decided_t, output_ta, mem_loss, nupd, oi, bc],
                  parallel_iterations=1, swap_memory=True)
              # pylint: enable=cell-var-from-loop
              gpu_out_idx.append(tf.squeeze(out_idx, [2]))
              outputs = output_ta.stack()
              outputs = tf.squeeze(outputs, [2, 3])  # Now l x b x nmaps
            else:
              # If beam_size is 0 or less, we don't have a decoder.
              mem_loss = 0.0
              outputs = tf.transpose(step[:, :, 1, :], [1, 0, 2])
              gpu_out_idx.append(tf.argmax(outputs, 2))
  
            # Final convolution to get logits, list outputs.
            outputs = tf.matmul(tf.reshape(outputs, [-1, nmaps]), output_w)
            outputs = tf.reshape(outputs, [length, batch_size, noclass])
          gpu_outputs[gpu] = tf.nn.softmax(outputs)
  
          # Calculate cross-entropy loss and normalize it.
          targets_soft = make_dense(tf.squeeze(gpu_target[gpu], [1]),
                                    noclass, 0.1)
          targets_soft = tf.reshape(targets_soft, [-1, noclass])
          targets_hard = make_dense(tf.squeeze(gpu_target[gpu], [1]),
                                    noclass, 0.0)
          targets_hard = tf.reshape(targets_hard, [-1, noclass])
          output = tf.transpose(outputs, [1, 0, 2])
          xent_soft = tf.reshape(tf.nn.softmax_cross_entropy_with_logits(
              logits=tf.reshape(output, [-1, noclass]), labels=targets_soft),
                                 [batch_size, length])
          xent_hard = tf.reshape(tf.nn.softmax_cross_entropy_with_logits(
              logits=tf.reshape(output, [-1, noclass]), labels=targets_hard),
                                 [batch_size, length])
          low, high = 0.1 / float(noclass - 1), 0.9
          const = high * tf.log(high) + float(noclass - 1) * low * tf.log(low)
          weight_sum = tf.reduce_sum(weights) + 1e-20
          true_perp = tf.reduce_sum(xent_hard * weights) / weight_sum
          soft_loss = tf.reduce_sum(xent_soft * weights) / weight_sum
          perp_loss = soft_loss + const
          # Final loss: cross-entropy + shared parameter relaxation part + extra.
          mem_loss = 0.5 * tf.reduce_mean(mem_loss) / length_float
          total_loss = perp_loss + mem_loss
          gpu_losses[gpu].append(true_perp)
  
          # Gradients.
          if backward:
            data.print_out("Creating backward pass for the model.")
            grads = tf.gradients(
                total_loss, tf.trainable_variables(),
                colocate_gradients_with_ops=True)
            for g_i, g in enumerate(grads):
              if isinstance(g, tf.IndexedSlices):
                grads[g_i] = tf.convert_to_tensor(g)
            grads, norm = tf.clip_by_global_norm(grads, max_grad_norm)
            gpu_grad_norms[gpu].append(norm)
            for g in grads:
              if grad_noise_scale > 0.001:
                g += tf.truncated_normal(tf.shape(g)) * self.noise_param
            grads_list.append(grads)
          else:
            gpu_grad_norms[gpu].append(0.0)
          data.print_out("Created model for gpu %d in %.2f s."
                         % (gpu, time.time() - start_time))
  
      self.updates = []
      self.after_enc_step = tf.concat(axis=0, values=self.after_enc_step)  # Concat GPUs.
      if backward:
        tf.get_variable_scope()._reuse = False
        tf.get_variable_scope().set_caching_device(None)
        grads = [gpu_avg([grads_list[g][i] for g in xrange(num_gpus)])
                 for i in xrange(len(grads_list[0]))]
        update = adam_update(grads)
        self.updates.append(update)
      else:
        self.updates.append(tf.no_op())
  
      self.losses = [gpu_avg([gpu_losses[g][i] for g in xrange(num_gpus)])
                     for i in xrange(len(gpu_losses[0]))]
      self.out_idx = tf.concat(axis=0, values=gpu_out_idx)
      self.grad_norms = [gpu_avg([gpu_grad_norms[g][i] for g in xrange(num_gpus)])
                         for i in xrange(len(gpu_grad_norms[0]))]
      self.outputs = [tf.concat(axis=1, values=[gpu_outputs[g] for g in xrange(num_gpus)])]
      self.quantize_op = quantize_weights_op(512, 8)
      if backward:
        self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
  
    def step(self, sess, inp, target, do_backward_in, noise_param=None,
             beam_size=2, eos_id=2, eos_cost=0.0, update_mem=None, state=None):
      """Run a step of the network."""
      batch_size, height, length = inp.shape[0], inp.shape[1], inp.shape[2]
      do_backward = do_backward_in
      train_mode = True
      if do_backward_in is None:
        do_backward = False
        train_mode = False
      if update_mem is None:
        update_mem = do_backward
      feed_in = {}
      # print "    feeding sequences of length %d" % length
      if state is None:
        state = np.zeros([batch_size, length, height, self.nmaps])
      feed_in[self.prev_step.name] = state
      feed_in[self.length_tensor.name] = length
      feed_in[self.noise_param.name] = noise_param if noise_param else 0.0
      feed_in[self.do_training.name] = 1.0 if do_backward else 0.0
      feed_in[self.update_mem.name] = 1 if update_mem else 0
      if do_backward_in is False:
        feed_in[self.sampling.name] = 0.0
      index = 0  # We're dynamic now.
      feed_out = []
      if do_backward:
        feed_out.append(self.updates[index])
        feed_out.append(self.grad_norms[index])
      if train_mode:
        feed_out.append(self.losses[index])
      feed_in[self.input.name] = inp
      feed_in[self.target.name] = target
      feed_out.append(self.outputs[index])
      if train_mode:
        # Make a full-sequence training step with one call to session.run.
        res = sess.run([self.after_enc_step] + feed_out, feed_in)
        after_enc_state, res = res[0], res[1:]
      else:
        # Make a full-sequence decoding step with one call to session.run.
        feed_in[self.sampling.name] = 1.1  # Sample every time.
        res = sess.run([self.after_enc_step, self.out_idx] + feed_out, feed_in)
        after_enc_state, out_idx = res[0], res[1]
        res = [res[2][l] for l in xrange(length)]
        outputs = [out_idx[:, i] for i in xrange(length)]
        cost = [0.0 for _ in xrange(beam_size * batch_size)]
        seen_eos = [0 for _ in xrange(beam_size * batch_size)]
        for idx, logit in enumerate(res):
          best = outputs[idx]
          for b in xrange(batch_size):
            if seen_eos[b] > 1:
              cost[b] -= eos_cost
            else:
              cost[b] += np.log(logit[b][best[b]])
            if best[b] in [eos_id]:
              seen_eos[b] += 1
        res = [[-c for c in cost]] + outputs
      # Collect and output results.
      offset = 0
      norm = None
      if do_backward:
        offset = 2
        norm = res[1]
      if train_mode:
        outputs = res[offset + 1]
        outputs = [outputs[l] for l in xrange(length)]
      return res[offset], outputs, norm, after_enc_state
  