dgl.transform.segmented_knn_graph

dgl.transform.segmented_knn_graph(x, k, segs)[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.

The matrices are concatenated along the first axis, and are segmented by segs. Each block would be transformed into a separate graph. The graphs will be unioned.

Parameters:
  • x (Tensor) – The input tensor.
  • k (int) – The number of neighbors
  • segs (iterable of int) – Number of points of each point set. Must sum up to the number of rows in x.
Returns:

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

Return type:

DGLGraph