Given a data_activation dictionary that contains a np array with activations for each layer,
    find the knns in the training data.

def find_train_knns(self, data_activations):
    """
    Given a data_activation dictionary that contains a np array with activations for each layer,
    find the knns in the training data.
    """
    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