I haven’t been able to find it implemented anywhere “officially”, the version i wrote in Tensorflow just called “tf.argmax”. It’s completely possible that there was something wrong with that model though, because the performance wasn’t really up to scratch. Here’s the relevant decoder code from my implementation:
def decoder_loop(iteration, s, e, us, ue, h_c_state, h_m_state, loss):
    decoder_endpoint_input = tf.concat([us, ue], axis=-1)
    _, (h_c_state, h_m_state)  = decoder_lstm_cell(decoder_endpoint_input, (h_c_state, h_m_state), scope=scope)
    h = h_m_state
    alpha = hmn(U, h, us, ue, "decoder_start_hmn")
    s = tf.argmax(alpha, axis=-1)
    us = tf.gather_nd(U, [[i, tens] for (i, tens) in enumerate(tf.unstack(s))])
    beta = hmn(U, h, us, ue, "decoder_end_hmn")
    e = tf.argmax(beta, axis=-1)
    ue = tf.gather_nd(U, [[i, tens] for (i, tens) in enumerate(tf.unstack(e))])
    loss += tf.reduce_sum(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=alpha, labels=answer_starts))
    loss += tf.reduce_sum(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=beta, labels=answer_ends))
    return (tf.add(iteration, 1), s, e, us, ue, h_c_state, h_m_state, loss)

_, s, e, _, _, _, _, loss = tf.while_loop(looping_cond, decoder_loop, (iteration_0, s_0, e_0, us_0, ue_0, h_c_state_0, h_m_state_0, loss_0), parallel_iterations=1)
tf.summary.scalar("loss", loss)