code: def knn(
      queries: torch.Tensor,
      index: faiss.IndexFlatL2,
      keys_store: np.ndarray,
      vals_store: np.ndarray,
      n_tokens: int,
  ):
      """
      ## $k$-NN to get $p(w_t, c_t)$
  
      Here we refer to $f(\textcolor{yellowgreen}{c_t})$ as queries,
      $f(c_i)$ as keys and $w_i$ as values.
      """
  
      # Save shape of queries to reshape results
      queries_shape = queries.shape
  
      # Flatten the `batch` and `sequence` dimensions of queries
      queries = queries.view(-1, queries_shape[-1])
  
      # Find 10 nearest neighbors of $f(\textcolor{yellowgreen}{c_t})$ among $f(c_i)$.
      # `distance` is the distance given by FAISS and `idx`, $i$ is the index of it in `keys_store`.
      distance, idx = index.search(queries.numpy(), 10)
  
      # Get $f(c_i)$
      keys_found = queries.new_tensor(keys_store[idx])
      # Get $w_i$
      vals_found = torch.tensor(vals_store[idx]).squeeze(-1)
  
      # We are going to calculate the cosine similarity between normalized vectors
  
      # Normalize $f(c_i)$
      keys_found_n = keys_found / torch.sqrt(
          (keys_found**2).sum(-1, keepdims=True) + 1e-10
      )
      # Normalize $f(\textcolor{yellowgreen}{c_t})$
      queries_n = queries / torch.sqrt(
          (queries**2).sum(-1, keepdims=True) + 1e-10
      )
  
      # Get the dot-product, or cosine similarity
      dot_prod = (keys_found_n * queries_n.unsqueeze(1)).sum(-1)
  
      # Token-wise logits
      logits_token = dot_prod.new_zeros(queries.shape[0], n_tokens)
      # Scatter and accumulate token logits based on the nearest neighbors
      _ = logits_token.scatter_(
          dim=1, index=vals_found, src=dot_prod, reduce="add"
      )
  
      # Reshape the logits
      logits_token = logits_token.reshape(
          queries_shape[0], queries_shape[1], -1
      )
  
      return logits_token
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/knn/eval_knn.py
imports: from typing import Optional, List
  import faiss
  import numpy
  import torch
  from labml import monit, lab
  from labml.logger import inspect
  from labml_nn.transformers.knn.train_model import Configs
label_desc: Use image node_id 1 for calling a global function with example usage: knn(queries, index, keys_store, vals_store, n_tokens) and returns: logits_token