import tensorflow as tf

      start_top_log_probs, start_top_index = tf.nn.top_k(
          start_log_probs, k=FLAGS.start_n_top)
      start_index = tf.one_hot(start_top_index,
                               depth=seq_len, axis=-1, dtype=tf.float32)
      start_features = tf.einsum("lbh,bkl->bkh", output, start_index)
      end_input = tf.tile(output[:, :, None],
                          [1, 1, FLAGS.start_n_top, 1])
      start_features = tf.tile(start_features[None],
                               [seq_len, 1, 1, 1])
      end_input = tf.concat([end_input, start_features], axis=-1)
