class_docstr: Utility for nearest neighbor search. It would first try to use `faiss`_ (requiring separate pacakge installment)
  as the backend for better computational performance. Otherwise, `scikit-learn` would be used as it is pre-installed
  with gpytorch.
  
  :param int k: number of nearest neighbors
  :param int dim: dimensionality of data
  :param torch.Size batch_shape: batch shape for train data
  :param str preferred_nnlib: currently supports `faiss` and `scikit-learn` (default: faiss).
  :param torch.device device: device that the NN search will be performed on.
  
  Example:
      >>> train_x = torch.randn(10, 5)
      >>> nn_util = NNUtil(k=3, dim=train_x.size(-1), device=train_x.device)
      >>> nn_util.set_nn_idx(train_x)
      >>> test_x = torch.randn(2, 5)
      >>> test_nn_indices = nn_util.find_nn_idx(test_x) # finding 3 nearest neighbors for test_x
      >>> test_nn_indices = nn_util.find_nn_idx(test_x, k=2) # finding 2 nearest neighbors for test_x
      >>> sequential_nn_idx = nn_util.build_sequential_nn_idx(train_x) # build up sequential nearest neighbor
      >>>     # structure for train_x
  
  .. _faiss:
      https://github.com/facebookresearch/faiss
class_docstr_tok: ["Utility","for","nearest","neighbor","search",".","It","would","first","try","to","use","`","faiss","`","_","(","requiring","separate","pacakge","installment",")","as","the","backend","for","better","computational","performance",".","Otherwise",",","`","scikit-learn","`","would","be","used","as","it","is","pre-installed","with","gpytorch",".",":","param","int","k",":","number","of","nearest","neighbors",":","param","int","dim",":","dimensionality","of","data",":","param","torch.Size","batch_shape",":","batch","shape","for","train","data",":","param","str","preferred_nnlib",":","currently","supports","`","faiss","`","and","`","scikit-learn","`","(","default",":","faiss",")",".",":","param","torch.device","device",":","device","that","the","NN","search","will","be","performed","on",".","Example",":",">",">",">","train_x","=","torch.randn","(","10",",","5",")",">",">",">","nn_util","=","NNUtil","(","k=3",",","dim=train_x.size","(","-1",")",",","device=train_x.device",")",">",">",">","nn_util.set_nn_idx","(","train_x",")",">",">",">","test_x","=","torch.randn","(","2",",","5",")",">",">",">","test_nn_indices","=","nn_util.find_nn_idx","(","test_x",")","#","finding","3","nearest","neighbors","for","test_x",">",">",">","test_nn_indices","=","nn_util.find_nn_idx","(","test_x",",","k=2",")","#","finding","2","nearest","neighbors","for","test_x",">",">",">","sequential_nn_idx","=","nn_util.build_sequential_nn_idx","(","train_x",")","#","build","up","sequential","nearest","neighbor",">",">",">","#","structure","for","train_x","..","_faiss",":","https",":","\/\/github.com\/facebookresearch\/faiss"]
docstr: Set the indices of training data to facilitate nearest neighbor search.
  This function needs to be called every time that the data changes.
  
  :param torch.Tensor train_x: training data points (... x N x D)
docstr_tok: ["Set","the","indices","of","training","data","to","facilitate","nearest","neighbor","search",".","This","function","needs","to","be","called","every","time","that","the","data","changes",".",":","param","torch.Tensor","train_x",":","training","data","points","(","...","x","N","x","D",")"]
code: def set_nn_idx(self, train_x):
      """
      Set the indices of training data to facilitate nearest neighbor search.
      This function needs to be called every time that the data changes.
  
      :param torch.Tensor train_x: training data points (... x N x D)
      """
      train_x = self._expand_and_check_shape(train_x)
      self.train_n = train_x.shape[-2]
  
      with torch.no_grad():
          if self.nnlib == "sklearn":
              self.train_neighbors = []
  
              from sklearn.neighbors import NearestNeighbors
  
              train_x = train_x.view(-1, self.train_n, self.dim)
  
              for i in range(self.batch_shape.numel()):
                  x = train_x[i].cpu().numpy()
                  self.train_neighbors.append(
                      NearestNeighbors(
                          n_neighbors=self.k, algorithm="auto"
                      ).fit(x)
                  )
          elif self.nnlib == "faiss":
              train_x = train_x.view(-1, self.train_n, self.dim)
              for i in range(self.batch_shape.numel()):
                  self.index[i].reset()
                  self.index[i].add(train_x[i])
  
file_path: gpytorch/gpytorch/utils/nearest_neighbors.py
imports: import warnings
  import torch
  from torch.nn import Module
label_desc: Use image node_id 5 for calling the NNUtil obj's underlying member method code with example usage: obj.set_nn_idx(train_x) without return types