Реализация лучевого поиска для нетекстовых задач в тензорном потоке

Question

Скажем, у меня есть ячейка РНН, обученная с помощью dynamic_rnn:

rnn_cell = tf.contrib.rnn.BasicLSTMCell(num_units)

# input_tensor: [batch_size, max_time, feature_size], each element of the last dimension is an sample
# rnn outputs: [batch_size, max_time, num_units]
# final_state: [batch_size, num_units]
rnn_outputs, final_state = tf.nn.dynamic_rnn(
        rnn_cell, input_tensor, dtype=tf.float32, 
        sequence_length=seq_length, time_major=False)

logits = tf.layers.dense(
        final_state.h, 2, 
        activation=None, 
        use_bias=True, 
        name='logits') # batch_size * 2

loss = tf.reduce_mean(
    tf.losses.sparse_softmax_cross_entropy(
        logits=logits, # batch_size * 2
        labels=labels))

# consider prob of class '1' as a score
score = tf.nn.softmax(logits)[:,1]

Теперь, учитывая input_tensor для прогнозирования, я хочу выполнить поиск луча, чтобы найти лучший ранг входных выборок, чтобы максимально увеличить общую оценку ранжированных выборок.