pyg_lib.ops

grouped_matmul(inputs: List[Tensor], others: List[Tensor], biases: Optional[List[Tensor]] = None) → List[Tensor][source]

Performs dense-dense matrix multiplication according to groups, utilizing dedicated kernels that effectively parallelize over groups.

inputs = [torch.randn(5, 16), torch.randn(3, 32)]
others = [torch.randn(16, 32), torch.randn(32, 64)]

outs = pyg_lib.ops.grouped_matmul(inputs, others)
assert len(outs) == 2
assert outs[0].size() == (5, 32)
assert outs[0] == inputs[0] @ others[0]
assert outs[1].size() == (3, 64)
assert outs[1] == inputs[1] @ others[1]

Parameters:

inputs (List[Tensor]) – List of left operand 2D matrices of shapes [N_i, K_i].
others (List[Tensor]) – List of right operand 2D matrices of shapes [K_i, M_i].
biases (Optional[List[Tensor]], default: None) – Optional bias terms to apply for each element.

Returns:

List[Tensor] – List of 2D output matrices of shapes [N_i, M_i].

segment_matmul(inputs: Tensor, ptr: Tensor, other: Tensor, bias: Optional[Tensor] = None) → Tensor[source]

Performs dense-dense matrix multiplication according to segments along the first dimension of inputs as given by ptr, utilizing dedicated kernels that effectively parallelize over groups.

inputs = torch.randn(8, 16)
ptr = torch.tensor([0, 5, 8])
other = torch.randn(2, 16, 32)

out = pyg_lib.ops.segment_matmul(inputs, ptr, other)
assert out.size() == (8, 32)
assert out[0:5] == inputs[0:5] @ other[0]
assert out[5:8] == inputs[5:8] @ other[1]

Parameters:

inputs (Tensor) – The left operand 2D matrix of shape [N, K].
ptr (Tensor) – Compressed vector of shape [B + 1], holding the boundaries of segments. For best performance, given as a CPU tensor.
other (Tensor) – The right operand 3D tensor of shape [B, K, M].
bias (Optional[Tensor], default: None) – The bias term of shape [B, M].

Returns:

Tensor – The 2D output matrix of shape [N, M].

sampled_add(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) → Tensor[source]

Performs a sampled addition of left and right according to the indices specified in left_index and right_index.

\[\textrm{out} = \textrm{left}[\textrm{left_index}] + \textrm{right}[\textrm{right_index}]\]

This operation fuses the indexing and addition operation together, thus being more runtime and memory-efficient.

Parameters:

left (Tensor) – The left tensor.
right (Tensor) – The right tensor.
left_index (Optional[Tensor], default: None) – The values to sample from the left tensor.
right_index (Optional[Tensor], default: None) – The values to sample from the right tensor.

Returns:

Tensor – The output tensor.

sampled_sub(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) → Tensor[source]

Performs a sampled subtraction of left by right according to the indices specified in left_index and right_index.

\[\textrm{out} = \textrm{left}[\textrm{left_index}] - \textrm{right}[\textrm{right_index}]\]

This operation fuses the indexing and subtraction operation together, thus being more runtime and memory-efficient.

Parameters:

left (Tensor) – The left tensor.
right (Tensor) – The right tensor.
left_index (Optional[Tensor], default: None) – The values to sample from the left tensor.
right_index (Optional[Tensor], default: None) – The values to sample from the right tensor.

Returns:

Tensor – The output tensor.

sampled_mul(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) → Tensor[source]

Performs a sampled multiplication of left and right according to the indices specified in left_index and right_index.

\[\textrm{out} = \textrm{left}[\textrm{left_index}] * \textrm{right}[\textrm{right_index}]\]

This operation fuses the indexing and multiplication operation together, thus being more runtime and memory-efficient.

Parameters:

left (Tensor) – The left tensor.
right (Tensor) – The right tensor.
left_index (Optional[Tensor], default: None) – The values to sample from the left tensor.
right_index (Optional[Tensor], default: None) – The values to sample from the right tensor.

Returns:

Tensor – The output tensor.

sampled_div(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) → Tensor[source]

Performs a sampled division of left by right according to the indices specified in left_index and right_index.

\[\textrm{out} = \textrm{left}[\textrm{left_index}] / \textrm{right}[\textrm{right_index}]\]

This operation fuses the indexing and division operation together, thus being more runtime and memory-efficient.

Parameters:

left (Tensor) – The left tensor.
right (Tensor) – The right tensor.
left_index (Optional[Tensor], default: None) – The values to sample from the left tensor.
right_index (Optional[Tensor], default: None) – The values to sample from the right tensor.

Returns:

Tensor – The output tensor.

index_sort(inputs: Tensor, max_value: Optional[int] = None) → Tuple[Tensor, Tensor][source]

Sorts the elements of the inputs tensor in ascending order. It is expected that inputs is one-dimensional and that it only contains positive integer values. If max_value is given, it can be used by the underlying algorithm for better performance.

Note

This operation is optimized only for tensors associated with the CPU device.

Parameters:

inputs (Tensor) – A vector with positive integer values.
max_value (Optional[int], default: None) – The maximum value stored inside inputs. This value can be an estimation, but needs to be greater than or equal to the real maximum.

Returns:

Tuple[Tensor, Tensor] – A tuple containing sorted values and indices of the elements in the original input tensor.

softmax_csr(src: Tensor, ptr: Tensor, dim: int = 0) → Tensor[source]

Computes a sparsely evaluated softmax. Given a value tensor src, this function first groups the values along the given dimension dim, based on the indices specified via ptr, and then proceeds to compute the softmax individually for each group.

Examples

>>> src = torch.randn(4, 4)
>>> ptr = torch.tensor([0, 4])
>>> softmax(src, ptr)
tensor([[0.0157, 0.0984, 0.1250, 0.4523],
        [0.1453, 0.2591, 0.5907, 0.2410],
        [0.0598, 0.2923, 0.1206, 0.0921],
        [0.7792, 0.3502, 0.1638, 0.2145]])

Parameters:

src (Tensor) – The source tensor.
ptr (Tensor) – Groups defined by CSR representation.
dim (int, default: 0) – The dimension in which to normalize.

Return type:

Tensor

scatter_sum(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor[source]

Reduces all values from the src tensor into out at the indices specified in the index tensor along a given axis dim, using sum as the reduction.

If out is not given, a new tensor is allocated and zero-initialized. If out is given, values are accumulated into it (no zero-init).

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The indices of elements to scatter.
dim (int, default: -1) – The axis along which to index.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, automatically create output with size dim_size at dimension dim. If dim_size is also None, a minimal sized output is returned.

Returns:

Tensor – The reduced tensor.

scatter_add(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor

Reduces all values from the src tensor into out at the indices specified in the index tensor along a given axis dim, using sum as the reduction.

If out is not given, a new tensor is allocated and zero-initialized. If out is given, values are accumulated into it (no zero-init).

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The indices of elements to scatter.
dim (int, default: -1) – The axis along which to index.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, automatically create output with size dim_size at dimension dim. If dim_size is also None, a minimal sized output is returned.

Returns:

Tensor – The reduced tensor.

scatter_mul(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor[source]

Reduces all values from the src tensor into out at the indices specified in the index tensor along a given axis dim, using mul as the reduction.

If out is not given, a new tensor is allocated and initialized to ones (the multiplicative identity). If out is given, values are multiplied into it (no ones-init).

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The indices of elements to scatter.
dim (int, default: -1) – The axis along which to index.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, automatically create output with size dim_size at dimension dim.

Returns:

Tensor – The reduced tensor.

scatter_mean(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor[source]

Reduces all values from the src tensor into out at the indices specified in the index tensor along a given axis dim, using mean as the reduction. Empty buckets yield zero.

For floating-point inputs, division is exact. For integer inputs, floor-division is used (matching upstream pytorch_scatter).

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The indices of elements to scatter.
dim (int, default: -1) – The axis along which to index.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, automatically create output with size dim_size at dimension dim.

Returns:

Tensor – The reduced tensor.

scatter_min(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tuple[Tensor, Tensor][source]

Reduces all values from the src tensor into out at the indices specified in the index tensor along a given axis dim, using min as the reduction.

Returns a tuple (values, argindex) where argindex[i] is the index along dim of the source element that contributed to values[i]. Empty buckets yield value 0 and argindex src.size(dim) (sentinel). Only values is differentiable.

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The indices of elements to scatter.
dim (int, default: -1) – The axis along which to index.
out (Optional[Tensor], default: None) – The destination tensor for the values.
dim_size (Optional[int], default: None) – If out is not given, automatically create output with size dim_size at dimension dim.

Returns:

Tuple[Tensor, Tensor] – Tuple (values, argindex).

scatter_max(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tuple[Tensor, Tensor][source]

Reduces all values from the src tensor into out at the indices specified in the index tensor along a given axis dim, using max as the reduction.

Returns a tuple (values, argindex) where argindex[i] is the index along dim of the source element that contributed to values[i]. Empty buckets yield value 0 and argindex src.size(dim) (sentinel). Only values is differentiable.

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The indices of elements to scatter.
dim (int, default: -1) – The axis along which to index.
out (Optional[Tensor], default: None) – The destination tensor for the values.
dim_size (Optional[int], default: None) – If out is not given, automatically create output with size dim_size at dimension dim.

Returns:

Tuple[Tensor, Tensor] – Tuple (values, argindex).

segment_sum_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor[source]

Reduces all values from the src tensor into out according to a sorted COO index, using sum as the reduction.

Unlike scatter_sum(), the reduction axis is always index.dim() - 1 (so no dim argument is exposed). The index must be sorted in ascending order along that axis.

If out is not given, a new zero-initialized tensor is allocated. If out is given, values are accumulated into it (no zero-init).

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – Sorted COO indices.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, the output size along the reduction axis. Auto-inferred from index.max() + 1 if None.

Returns:

Tensor – The reduced tensor.

segment_add_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor

Reduces all values from the src tensor into out according to a sorted COO index, using sum as the reduction.

Unlike scatter_sum(), the reduction axis is always index.dim() - 1 (so no dim argument is exposed). The index must be sorted in ascending order along that axis.

If out is not given, a new zero-initialized tensor is allocated. If out is given, values are accumulated into it (no zero-init).

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – Sorted COO indices.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, the output size along the reduction axis. Auto-inferred from index.max() + 1 if None.

Returns:

Tensor – The reduced tensor.

segment_mean_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tensor[source]

Reduces all values from the src tensor into out according to a sorted COO index, using mean as the reduction. Empty buckets yield zero.

The reduction axis is index.dim() - 1. For integer dtypes, floor-division is used.

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – Sorted COO indices.
out (Optional[Tensor], default: None) – The destination tensor.
dim_size (Optional[int], default: None) – If out is not given, the output size along the reduction axis. Auto-inferred from index.max() + 1 if None.

Returns:

Tensor – The reduced tensor.

segment_min_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tuple[Tensor, Tensor][source]

Reduces all values from the src tensor into out according to a sorted COO index, using min as the reduction.

Returns a tuple (values, argindex). The reduction axis is index.dim() - 1. Empty buckets yield value 0 and argindex src.size(dim) (sentinel). Only values is differentiable.

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – Sorted COO indices.
out (Optional[Tensor], default: None) – The destination tensor for the values.
dim_size (Optional[int], default: None) – If out is not given, the output size along the reduction axis.

Returns:

Tuple[Tensor, Tensor] – Tuple (values, argindex).

segment_max_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None, dim_size: Optional[int] = None) → Tuple[Tensor, Tensor][source]

Reduces all values from the src tensor into out according to a sorted COO index, using max as the reduction.

Returns a tuple (values, argindex). The reduction axis is index.dim() - 1. Empty buckets yield value 0 and argindex src.size(dim) (sentinel). Only values is differentiable.

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – Sorted COO indices.
out (Optional[Tensor], default: None) – The destination tensor for the values.
dim_size (Optional[int], default: None) – If out is not given, the output size along the reduction axis.

Returns:

Tuple[Tensor, Tensor] – Tuple (values, argindex).

gather_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None) → Tensor[source]

Gathers values from src at positions specified by index along axis index.dim() - 1. This is the symmetric inverse of segment_sum_coo().

Parameters:

src (Tensor) – The source tensor.
index (Tensor) – The COO indices to gather.
out (Optional[Tensor], default: None) – The destination tensor (overwritten if given).

Returns:

Tensor – out[i] = src[index[i]] along the gather axis.

segment_sum_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None) → Tensor[source]

Reduces all values from the src tensor into out according to a CSR indptr, using sum as the reduction.

The reduction axis is indptr.dim() - 1 and the output size along that axis is indptr.size(-1) - 1 (no explicit dim_size arg).

If out is not given, a new zero-initialized tensor is allocated. If out is given, values are accumulated into it.

Parameters:

src (Tensor) – The source tensor.
indptr (Tensor) – The CSR row pointer of shape [..., R+1].
out (Optional[Tensor], default: None) – The destination tensor.

Returns:

Tensor – The reduced tensor.

segment_add_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None) → Tensor

Reduces all values from the src tensor into out according to a CSR indptr, using sum as the reduction.

The reduction axis is indptr.dim() - 1 and the output size along that axis is indptr.size(-1) - 1 (no explicit dim_size arg).

If out is not given, a new zero-initialized tensor is allocated. If out is given, values are accumulated into it.

Parameters:

src (Tensor) – The source tensor.
indptr (Tensor) – The CSR row pointer of shape [..., R+1].
out (Optional[Tensor], default: None) – The destination tensor.

Returns:

Tensor – The reduced tensor.

segment_mean_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None) → Tensor[source]

Reduces all values from the src tensor into out according to a CSR indptr, using mean as the reduction. Empty rows yield zero.

Parameters:

src (Tensor) – The source tensor.
indptr (Tensor) – The CSR row pointer of shape [..., R+1].
out (Optional[Tensor], default: None) – The destination tensor.

Returns:

Tensor – The reduced tensor.

segment_min_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None) → Tuple[Tensor, Tensor][source]

Reduces all values from the src tensor into out according to a CSR indptr, using min as the reduction.

Returns a tuple (values, argindex). Empty rows yield value 0 and argindex src.size(dim) (sentinel). Only values is differentiable.

Parameters:

src (Tensor) – The source tensor.
indptr (Tensor) – The CSR row pointer of shape [..., R+1].
out (Optional[Tensor], default: None) – The destination tensor for the values.

Returns:

Tuple[Tensor, Tensor] – Tuple (values, argindex).

segment_max_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None) → Tuple[Tensor, Tensor][source]

Reduces all values from the src tensor into out according to a CSR indptr, using max as the reduction.

Returns a tuple (values, argindex). Empty rows yield value 0 and argindex src.size(dim) (sentinel). Only values is differentiable.

Parameters:

src (Tensor) – The source tensor.
indptr (Tensor) – The CSR row pointer of shape [..., R+1].
out (Optional[Tensor], default: None) – The destination tensor for the values.

Returns:

Tuple[Tensor, Tensor] – Tuple (values, argindex).

gather_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None) → Tensor[source]

Gathers values from src to all positions encoded by the CSR indptr: for each row r, broadcasts src[r] to all output positions in [indptr[r], indptr[r+1]). Symmetric inverse of segment_sum_csr().

Parameters:

src (Tensor) – The source tensor.
indptr (Tensor) – The CSR row pointer of shape [..., R+1].
out (Optional[Tensor], default: None) – The destination tensor (overwritten if given).

Returns:

Tensor – Broadcast tensor of shape src.shape with indptr.dim()-1 axis replaced by indptr[-1].

scatter(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None, reduce: str = 'sum') → Tensor[source]

Polymorphic scatter dispatcher.

Routes to the typed scatter op based on reduce: "sum"/"add" -> scatter_sum(), "mul" -> scatter_mul(), "mean" -> scatter_mean(), "min" -> scatter_min(), "max" -> scatter_max(). Min/max return only the values.

Return type:: Tensor

segment_coo(src: Tensor, index: Tensor, out: Optional[Tensor] = None, dim_size: Optional[int] = None, reduce: str = 'sum') → Tensor[source]

Polymorphic COO segment dispatcher.

Routes by reduce to the typed segment_*_coo op. Min/max return only the value tensor.

Return type:: Tensor

segment_csr(src: Tensor, indptr: Tensor, out: Optional[Tensor] = None, reduce: str = 'sum') → Tensor[source]

Polymorphic CSR segment dispatcher.

Routes by reduce to the typed segment_*_csr op. Min/max return only the value tensor.

Return type:: Tensor

scatter_softmax(src: Tensor, index: Tensor, dim: int = -1, dim_size: Optional[int] = None) → Tensor[source]

Per-bucket softmax composite. Float inputs only.

Return type:: Tensor

scatter_log_softmax(src: Tensor, index: Tensor, dim: int = -1, dim_size: Optional[int] = None, eps: float = 1e-12) → Tensor[source]

Per-bucket log-softmax composite. Float inputs only.

Return type:: Tensor

scatter_std(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None, unbiased: bool = True) → Tensor[source]

Per-bucket standard deviation composite. Float inputs only.

Ports torch_scatter.composite.scatter_std: uses scatter_sum() for both passes (avoids integer floor-division coupling that would appear if scatter_mean() were used). Applies Bessel’s correction N / (N - 1) when unbiased is True.

Return type:: Tensor

scatter_logsumexp(src: Tensor, index: Tensor, dim: int = -1, out: Optional[Tensor] = None, dim_size: Optional[int] = None, eps: float = 1e-12) → Tensor[source]

Per-bucket log-sum-exp composite. Float inputs only.

Numerically stable: recenter by the per-bucket max before exponentiating. When out is supplied, positions that ended up non-finite (empty buckets that produced -inf) are restored from the caller-supplied out to preserve any caller-provided initial values.

Return type:: Tensor

spline_basis(pseudo: Tensor, kernel_size: Tensor, is_open_spline: Tensor, degree: int = 1) → Tuple[Tensor, Tensor][source]

Computes the B-spline basis functions.

Parameters:

pseudo (Tensor) – Pseudo-coordinates of shape [E, D].
kernel_size (Tensor) – Kernel size in each dimension of shape [D].
is_open_spline (Tensor) – Whether to use open B-splines of shape [D].
degree (int, default: 1) – B-spline degree (1, 2, or 3).

Returns:

Tuple[Tensor, Tensor] – Basis values of shape [E, S] and weight indices of shape [E, S].

spline_weighting(x: Tensor, weight: Tensor, basis: Tensor, weight_index: Tensor) → Tensor[source]

Computes the spline weighting of input features.

Parameters:

x (Tensor) – Input features of shape [E, M_in].
weight (Tensor) – Weight tensor of shape [K, M_in, M_out].
basis (Tensor) – B-spline basis values of shape [E, S].
weight_index (Tensor) – Weight indices of shape [E, S].

Returns:

Tensor – Output features of shape [E, M_out].

grid_cluster(pos: Tensor, size: Tensor, start: Optional[Tensor] = None, end: Optional[Tensor] = None) → Tensor[source]

Clusters all points in pos into voxels of size size.

Each point is assigned a cluster index based on which voxel it falls into. The voxel grid is defined by the size parameter and optionally bounded by start and end.

Parameters:

pos (Tensor) – Point positions of shape [N, D].
size (Tensor) – Voxel size in each dimension of shape [D].
start (Optional[Tensor], default: None) – Start of the voxel grid in each dimension of shape [D]. If None, uses the minimum of pos.
end (Optional[Tensor], default: None) – End of the voxel grid in each dimension of shape [D]. If None, uses the maximum of pos.

Returns:

Tensor – Cluster index for each point of shape [N].

fps(src: Tensor, ptr: Tensor, ratio: float = 0.5, random_start: bool = True) → Tensor[source]

Performs greedy farthest point sampling.

Starting from a random point (or the first point), iteratively selects the point that is farthest from the already selected set.

Parameters:

src (Tensor) – Point positions of shape [N, D].
ptr (Tensor) – Batch boundaries as a CSR pointer of shape [B + 1].
ratio (float, default: 0.5) – Fraction of points to sample from each batch (in (0, 1]).
random_start (bool, default: True) – If True, starts from a random point.

Returns:

Tensor – Indices of the sampled points of shape [M].

knn(x: Tensor, y: Tensor, k: int = 1, ptr_x: Optional[Tensor] = None, ptr_y: Optional[Tensor] = None, cosine: bool = False, num_workers: int = 1) → Tensor[source]

Finds for each element in y the k nearest points in x.

Parameters:

x (Tensor) – Reference points of shape [N, D].
y (Tensor) – Query points of shape [M, D].
k (int, default: 1) – Number of nearest neighbors.
ptr_x (Optional[Tensor], default: None) – Batch boundaries for x as a CSR pointer.
ptr_y (Optional[Tensor], default: None) – Batch boundaries for y as a CSR pointer.
cosine (bool, default: False) – If True, uses cosine distance (CUDA only).
num_workers (int, default: 1) – Number of workers (unused, for API compat).

Returns:

Tensor – Edge indices of shape [2, M*k] where row 0 is query indices and row 1 is reference indices.

radius(x: Tensor, y: Tensor, r: float = 1.0, ptr_x: Optional[Tensor] = None, ptr_y: Optional[Tensor] = None, max_num_neighbors: int = 32, num_workers: int = 1, ignore_same_index: bool = False) → Tensor[source]

Finds all points in x within distance r of points in y.

Parameters:

x (Tensor) – Reference points of shape [N, D].
y (Tensor) – Query points of shape [M, D].
r (float, default: 1.0) – Radius.
ptr_x (Optional[Tensor], default: None) – Batch boundaries for x as a CSR pointer.
ptr_y (Optional[Tensor], default: None) – Batch boundaries for y as a CSR pointer.
max_num_neighbors (int, default: 32) – Maximum number of neighbors per query point.
num_workers (int, default: 1) – Number of workers (unused, for API compat).
ignore_same_index (bool, default: False) – If True, ignores pairs with same index.

Returns:

Tensor – Edge indices of shape [2, E] where row 0 is query indices and row 1 is reference indices.

nearest(x: Tensor, y: Tensor, ptr_x: Optional[Tensor] = None, ptr_y: Optional[Tensor] = None) → Tensor[source]

Finds the nearest point in y for each point in x.

Parameters:

x (Tensor) – Query points of shape [N, D].
y (Tensor) – Reference points of shape [M, D].
ptr_x (Optional[Tensor], default: None) – Batch boundaries for x as a CSR pointer.
ptr_y (Optional[Tensor], default: None) – Batch boundaries for y as a CSR pointer.

Returns:

Tensor – Index tensor of shape [N] with the index of the nearest point in y for each point in x.

graclus_cluster(rowptr: Tensor, col: Tensor, weight: Optional[Tensor] = None) → Tensor[source]

Computes a greedy graph clustering via the Graclus algorithm.

Nodes are matched greedily in random order. The cluster ID for a matched pair (u, v) is min(u, v). Unmatched nodes are assigned their own index as cluster ID.

Parameters:

rowptr (Tensor) – CSR row pointer of shape [N + 1].
col (Tensor) – Column indices of shape [E].
weight (Optional[Tensor], default: None) – Optional edge weights of shape [E].

Returns:

Tensor – Cluster assignment of shape [N].

edge_sample(start: Tensor, rowptr: Tensor, count: int = 0, factor: float = 1.0) → Tensor[source]

Samples edges incident to the given start nodes.

For each start node, samples up to count edges. If count < 1, samples ceil(factor * degree) edges instead.

Parameters:

start (Tensor) – Start node indices of shape [S].
rowptr (Tensor) – CSR row pointer of shape [N + 1].
count (int, default: 0) – Fixed number of edges to sample per node. If < 1, uses factor instead.
factor (float, default: 1.0) – Fraction of edges to sample when count < 1.

Returns:

Tensor – Sampled edge indices (into the edge list).