dgl.transform.knn_graph

dgl.transform.knn_graph(x, k)[source]

Transforms the given point set to a directed graph, whose coordinates are given as a matrix. The predecessors of each point are its k-nearest neighbors.

If a 3D tensor is given instead, then each row would be transformed into a separate graph. The graphs will be unioned.

Parameters:
  • x (Tensor) –

    The input tensor.

    If 2D, each row of x corresponds to a node.

    If 3D, a k-NN graph would be constructed for each row. Then the graphs are unioned.

  • k (int) – The number of neighbors
Returns:

The graph. The node IDs are in the same order as x.

Return type:

DGLGraph