torch_geometric.nn.pool.approx_knn_graph

approx_knn_graph(x: Tensor, k: int, batch: Optional[Tensor] = None, loop: bool = False, flow: str = 'source_to_target') → Tensor[source]

Computes graph edges to the nearest approximated k points.

Note

Approximated \(k\)-nearest neighbor search is performed via the pynndescent library.

Parameters:

  • x (torch.Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{N \times F}\).

  • k (int) – The number of neighbors.

  • batch (torch.Tensor, optional) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example. (default: None)

  • loop (bool, optional) – If True, the graph will contain self-loops. (default: False)

  • flow (str, optional) – The flow direction when using in combination with message passing ("source_to_target" or "target_to_source"). (default: "source_to_target")

Return type:

torch.Tensor