import tensorflow as tf

            logits = generate(output, input_, context)

            pos = tf.expand_dims(pos, axis=1)
            state = tf.concat([state, context, pos, new_weights], axis=1)
            return state, logits

