'Queries memory for nearest neighbor.
Args:
query_vec: A batch of vectors to query (embedding of input to model).
intended_output: The values that would be the correct output of the
memory.
use_recent_idx: Whether to always insert at least one instance of a
correct memory fetch.
Returns:
A tuple (result, mask, teacher_loss).
result: The result of the memory look up.
mask: The affinity of the query to the result.
teacher_loss: The loss for training the memory module.'

def query(self, query_vec, intended_output, use_recent_idx=True):

    batch_size = tf.shape(query_vec)[0]
    output_given = (intended_output is not None)
    query_vec = tf.matmul(query_vec, self.query_proj)
    normalized_query = tf.nn.l2_normalize(query_vec, dim=1)
    hint_pool_idxs = self.get_hint_pool_idxs(normalized_query)
    if (output_given and use_recent_idx):
        most_recent_hint_idx = tf.gather(self.recent_idx, intended_output)
        hint_pool_idxs = tf.concat(axis=1, values=[hint_pool_idxs, tf.expand_dims(most_recent_hint_idx, 1)])
    choose_k = tf.shape(hint_pool_idxs)[1]
    with tf.device(self.var_cache_device):
        my_mem_keys = tf.stop_gradient(tf.gather(self.mem_keys, hint_pool_idxs, name='my_mem_keys_gather'))
        similarities = tf.matmul(tf.expand_dims(normalized_query, 1), my_mem_keys, adjoint_b=True, name='batch_mmul')
        hint_pool_sims = tf.squeeze(similarities, [1], name='hint_pool_sims')
        hint_pool_mem_vals = tf.gather(self.mem_vals, hint_pool_idxs, name='hint_pool_mem_vals')
    softmax_temp = max(1.0, (np.log((0.2 * self.choose_k)) / self.alpha))
    mask = tf.nn.softmax((hint_pool_sims[:, :(choose_k - 1)] * softmax_temp))
    teacher_hints = tf.to_float(tf.abs((tf.expand_dims(intended_output, 1) - hint_pool_mem_vals)))
    teacher_hints = (1.0 - tf.minimum(1.0, teacher_hints))
    (teacher_vals, teacher_hint_idxs) = tf.nn.top_k((hint_pool_sims * teacher_hints), k=1)
    (neg_teacher_vals, _) = tf.nn.top_k((hint_pool_sims * (1 - teacher_hints)), k=1)
    teacher_idxs = tf.gather(tf.reshape(hint_pool_idxs, [(-1)]), (teacher_hint_idxs[:, 0] + (choose_k * tf.range(batch_size))))
    teacher_vals *= (1 - tf.to_float(tf.equal(0.0, tf.reduce_sum(teacher_hints, 1))))
    nearest_neighbor = tf.to_int32(tf.argmax(hint_pool_sims[:, :(choose_k - 1)], 1))
    no_teacher_idxs = tf.gather(tf.reshape(hint_pool_idxs, [(-1)]), (nearest_neighbor + (choose_k * tf.range(batch_size))))
    sliced_hints = tf.slice(teacher_hints, [0, 0], [(-1), self.correct_in_top])
    incorrect_memory_lookup = tf.equal(0.0, tf.reduce_sum(sliced_hints, 1))
    teacher_loss = (tf.nn.relu(((neg_teacher_vals - teacher_vals) + self.alpha)) - self.alpha)
    with tf.device(self.var_cache_device):
        result = tf.gather(self.mem_vals, tf.reshape(no_teacher_idxs, [(-1)]))
    update_keys = normalized_query
    update_vals = intended_output
    fetched_idxs = teacher_idxs
    with tf.device(self.var_cache_device):
        fetched_keys = tf.gather(self.mem_keys, fetched_idxs, name='fetched_keys')
        fetched_vals = tf.gather(self.mem_vals, fetched_idxs, name='fetched_vals')
    fetched_keys_upd = (update_keys + fetched_keys)
    fetched_keys_upd = tf.nn.l2_normalize(fetched_keys_upd, dim=1)
    mem_age_with_noise = (self.mem_age + tf.random_uniform([self.memory_size], (- self.age_noise), self.age_noise))
    (_, oldest_idxs) = tf.nn.top_k(mem_age_with_noise, k=batch_size, sorted=False)
    with tf.control_dependencies([result]):
        upd_idxs = tf.where(incorrect_memory_lookup, oldest_idxs, fetched_idxs)
        upd_keys = tf.where(incorrect_memory_lookup, update_keys, fetched_keys_upd)
        upd_vals = tf.where(incorrect_memory_lookup, update_vals, fetched_vals)
    def make_update_op():
        return self.make_update_op(upd_idxs, upd_keys, upd_vals, batch_size, use_recent_idx, intended_output)
    update_op = tf.cond(self.update_memory, make_update_op, tf.no_op)
    with tf.control_dependencies([update_op]):
        result = tf.identity(result)
        mask = tf.identity(mask)
        teacher_loss = tf.identity(teacher_loss)
    return (result, mask, tf.reduce_mean(teacher_loss))
