'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):

    hash_slot_idxs = self.get_hash_slots(normalized_query)
    hint_pool_idxs = [tf.maximum(tf.minimum(tf.gather(self.hash_slots[i], idxs), (self.memory_size - 1)), 0) for (i, idxs) in enumerate(hash_slot_idxs)]
    return tf.concat(axis=1, values=hint_pool_idxs)
