radius(x: Tensor, y: Tensor, r: float, batch_x: = None, batch_y: = None, max_num_neighbors: int = 32, num_workers: int = 1) [source]

Finds for each element in y all points in x within distance r.

import torch

x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.Tensor([[-1, 0], [1, 0]])
batch_y = torch.tensor([0, 0])
assign_index = radius(x, y, 1.5, batch_x, batch_y)

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

• y (torch.Tensor) – Node feature matrix $$\mathbf{Y} \in \mathbb{R}^{M \times F}$$.

• r (float) – The radius.

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

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

• max_num_neighbors (int, optional) – The maximum number of neighbors to return for each element in y. (default: 32)

• num_workers (int, optional) – Number of workers to use for computation. Has no effect in case batch_x or batch_y is not None, or the input lies on the GPU. (default: 1)

Return type

torch.Tensor