code: def diagonal_neural_gpu(inputs, hparams, name=None):
    
    with tf.variable_scope(name, "diagonal_neural_gpu"):
  
      def step(state_tup, inp):
        """Single step of the improved Neural GPU."""
        state, _ = state_tup
        x = state
        for layer in range(hparams.num_hidden_layers):
          x, new_loss = common_layers.diagonal_conv_gru(
              x, (hparams.kernel_height, hparams.kernel_width),
              hparams.hidden_size,
              dropout=hparams.dropout,
              name="dcgru_%d" % layer)
        # Padding input is zeroed-out in the modality, we check this by summing.
        padding_inp = tf.less(tf.reduce_sum(tf.abs(inp), axis=[1, 2]), 0.00001)
        new_state = tf.where(padding_inp, state, x)  # No-op where inp is padding.
        return new_state, new_loss
  
      final_state, losses = tf.scan(
          step,
          tf.transpose(inputs, [1, 0, 2, 3]),
          initializer=(inputs, tf.constant(0.0)),
          parallel_iterations=1,
          swap_memory=True)
      return final_state[0, :, :, :, :], 2.0 * tf.reduce_mean(losses)
docs: Improved Neural GPU as in https://arxiv.org/abs/1702.08727.