'Get small set of idxs to compute nearest neighbor queries on.
This is an expensive look-up on the whole memory that is used to
avoid more expensive operations later on.
Args:
normalized_query: A Tensor of shape [None, key_dim].
Returns:
A Tensor of shape [None, choose_k] of indices in memory
that are closest to the queries.'

def get_hint_pool_idxs(self, normalized_query):

    with tf.device(self.nn_device):
        similarities = tf.matmul(tf.stop_gradient(normalized_query), self.mem_keys, transpose_b=True, name='nn_mmul')
    (_, hint_pool_idxs) = tf.nn.top_k(tf.stop_gradient(similarities), k=self.choose_k, name='nn_topk')
    return hint_pool_idxs
