class_docstr: Layer that transforms one point set into a graph, or a batch of
  point sets with the same number of points into a batched union of those graphs.
  
  The KNNGraph is implemented in the following steps:
  
  1. Compute an NxN matrix of pairwise distance for all points.
  2. Pick the k points with the smallest distance for each point as their k-nearest neighbors.
  3. Construct a graph with edges to each point as a node from its k-nearest neighbors.
  
  The overall computational complexity is :math:`O(N^2(logN + D)`.
  
  If a batch of point sets is provided, the point :math:`j` in point
  set :math:`i` is mapped to graph node ID: :math:`i \times M + j`, where
  :math:`M` is the number of nodes in each point set.
  
  The predecessors of each node are the k-nearest neighbors of the
  corresponding point.
  
  Parameters
  ----------
  k : int
      The number of neighbors.
  
  Notes
  -----
  The nearest neighbors found for a node include the node itself.
  
  Examples
  --------
  The following example uses PyTorch backend.
  
  >>> import torch
  >>> from dgl.nn.pytorch.factory import KNNGraph
  >>>
  >>> kg = KNNGraph(2)
  >>> x = torch.tensor([[0,1],
                        [1,2],
                        [1,3],
                        [100, 101],
                        [101, 102],
                        [50, 50]])
  >>> g = kg(x)
  >>> print(g.edges())
      (tensor([0, 1, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5]),
       tensor([0, 0, 1, 2, 1, 2, 5, 3, 4, 3, 4, 5]))
class_docstr_tok: ["Layer","that","transforms","one","point","set","into","a","graph",",","or","a","batch","of","point","sets","with","the","same","number","of","points","into","a","batched","union","of","those","graphs",".","The","KNNGraph","is","implemented","in","the","following","steps",":","1",".","Compute","an","NxN","matrix","of","pairwise","distance","for","all","points",".","2",".","Pick","the","k","points","with","the","smallest","distance","for","each","point","as","their","k-nearest","neighbors",".","3",".","Construct","a","graph","with","edges","to","each","point","as","a","node","from","its","k-nearest","neighbors",".","The","overall","computational","complexity","is",":","math",":","`","O","(","N^2","(","logN","+","D",")","`",".","If","a","batch","of","point","sets","is","provided",",","the","point",":","math",":","`","j","`","in","point","set",":","math",":","`","i","`","is","mapped","to","graph","node","ID",":",":","math",":","`","i","\\times","M","+","j","`",",","where",":","math",":","`","M","`","is","the","number","of","nodes","in","each","point","set",".","The","predecessors","of","each","node","are","the","k-nearest","neighbors","of","the","corresponding","point",".","Parameters","--","--","--","--","--","k",":","int","The","number","of","neighbors",".","Notes","--","--","-","The","nearest","neighbors","found","for","a","node","include","the","node","itself",".","Examples","--","--","--","--","The","following","example","uses","PyTorch","backend",".",">",">",">","import","torch",">",">",">","from","dgl.nn.pytorch.factory","import","KNNGraph",">",">",">",">",">",">","kg","=","KNNGraph","(","2",")",">",">",">","x","=","torch.tensor","(","[","[","0,1","]",",","[","1,2","]",",","[","1,3","]",",","[","100",",","101","]",",","[","101",",","102","]",",","[","50",",","50","]","]",")",">",">",">","g","=","kg","(","x",")",">",">",">","print","(","g.edges","(",")",")","(","tensor","(","[","0",",","1",",","1",",","1",",","2",",","2",",","2",",","3",",","3",",","4",",","4",",","5","]",")",",","tensor","(","[","0",",","0",",","1",",","2",",","1",",","2",",","5",",","3",",","4",",","3",",","4",",","5","]",")",")"]
code: def __init__(self, k):
      super(KNNGraph, self).__init__()
      self.k = k
  
file_path: dgl/python/dgl/nn/pytorch/factory.py
imports: import torch.nn
  from ...transforms import knn_graph, radius_graph, segmented_knn_graph
label_desc: Use image node_id 1 to create a new KNNGraph object from inherited base classes: nn with example: obj = KNNGraph(k)