import tensorflow as tf

    context_size = initial_context.shape[1].value

    def get_logits(state, ids, time):  # for beam-search decoding
        with tf.variable_scope('decoder_{}'.format(decoder.name)):
            state, context, pos, prev_weights = tf.split(state, [cell_state_size, context_size, 1, -1], axis=1)
            input_ = embed(ids)

            pos = tf.squeeze(pos, axis=1)
            pos = tf.cond(tf.equal(time, 0),
                          lambda: pos,
                          lambda: update_pos(pos, ids, encoder_input_length[align_encoder_id]))

            if decoder.cell_type.lower() == 'lstm' and decoder.use_lstm_full_state:
                output = state
            else:
                # Output is always the right-most part of the state (even with multi-layer RNNs)
