code: def find_train_knns(self, data_activations):
      
      knns_ind = {}
      knns_labels = {}
  
      for layer in self.layers:
        # Pre-process representations of data to normalize and remove training data mean.
        data_activations_layer = copy.copy(data_activations[layer])
        nb_data = data_activations_layer.shape[0]
        data_activations_layer /= np.linalg.norm(
            data_activations_layer, axis=1).reshape(-1, 1)
        data_activations_layer -= self.centers[layer]
  
        # Use FALCONN to find indices of nearest neighbors in training data.
        knns_ind[layer] = np.zeros(
            (data_activations_layer.shape[0], self.neighbors), dtype=np.int32)
        knn_errors = 0
        for i in range(data_activations_layer.shape[0]):
          query_res = self.query_objects[layer].find_k_nearest_neighbors(
              data_activations_layer[i], self.neighbors)
          try:
            knns_ind[layer][i, :] = query_res
          except:  # pylint: disable-msg=W0702
            knns_ind[layer][i, :len(query_res)] = query_res
            knn_errors += knns_ind[layer].shape[1] - len(query_res)
  
        # Find labels of neighbors found in the training data.
        knns_labels[layer] = np.zeros((nb_data, self.neighbors), dtype=np.int32)
        for data_id in range(nb_data):
          knns_labels[layer][data_id, :] = self.train_labels[knns_ind[layer][data_id]]
  
      return knns_ind, knns_labels
docs: Given a data_activation dictionary that contains a np array with activations for each layer,
      find the knns in the training data.