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: Find :math:`k` nearest neighbors for test data `test_x` among the training data stored in this utility
  
  :param test_x: test data, shape (... x N x D)
  :param int k: number of nearest neighbors. Default is the value used in utility initialization.
  :rtype: torch.LongTensor
  :return: the indices of nearest neighbors in the training data
docstr_tok: ["Find",":","math",":","`","k","`","nearest","neighbors","for","test","data","`","test_x","`","among","the","training","data","stored","in","this","utility",":","param","test_x",":","test","data",",","shape","(","...","x","N","x","D",")",":","param","int","k",":","number","of","nearest","neighbors",".","Default","is","the","value","used","in","utility","initialization",".",":","rtype",":","torch.LongTensor",":","return",":","the","indices","of","nearest","neighbors","in","the","training","data"]
code: def find_nn_idx(self, test_x, k=None):
      """
      Find :math:`k` nearest neighbors for test data `test_x` among the training data stored in this utility
  
      :param test_x: test data, shape (... x N x D)
      :param int k: number of nearest neighbors. Default is the value used in utility initialization.
      :rtype: torch.LongTensor
      :return: the indices of nearest neighbors in the training data
      """
  
      assert (
          self.train_n is not None
      ), "Please initialize with training data first."
      if k is None:
          k = self.k
      else:
          assert k > 0, f"k must be greater than 0, but got k = {k}."
      assert k <= self.train_n, (
          f"k should be smaller than number of train data, "
          f"but got k = {k}, number of train data = {self.train_n}."
      )
  
      test_x = self._expand_and_check_shape(test_x)
  
      test_n = test_x.shape[-2]
      test_x = test_x.view(-1, test_n, self.dim)
      nn_idx = torch.empty(
          self.batch_shape.numel(),
          test_n,
          k,
          dtype=torch.int64,
          device=test_x.device,
      )
  
      with torch.no_grad():
          if self.nnlib == "sklearn":
              if self.train_neighbors is None:
                  raise RuntimeError(
                      "The nearest neighbor set has not been defined. First call `set_nn_idx`"
                  )
  
              for i in range(self.batch_shape.numel()):
                  nn_idx_i = torch.from_numpy(
                      self.train_neighbors[i].kneighbors(
                          test_x[i].cpu().numpy()
                      )[1][..., :k]
                  )
                  nn_idx[i] = nn_idx_i.long().to(test_x.device)
          else:
              for i in range(self.batch_shape.numel()):
                  nn_idx[i] = self.index[i].search(test_x[i], k)[1]
  
      nn_idx = nn_idx.view(*self.batch_shape, test_n, k)
      return nn_idx
  
file_path: gpytorch/gpytorch/utils/nearest_neighbors.py
imports: import warnings
  import torch
  from torch.nn import Module
label_desc: Use image node_id 4 for calling the NNUtil obj's underlying member method code with example usage: obj.find_nn_idx(test_x, k) and returns: nn_idx