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]
- 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
inputsas given byptr, 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
leftandrightaccording to the indices specified inleft_indexandright_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.
- sampled_sub(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) Tensor[source]
Performs a sampled subtraction of
leftbyrightaccording to the indices specified inleft_indexandright_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.
- sampled_mul(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) Tensor[source]
Performs a sampled multiplication of
leftandrightaccording to the indices specified inleft_indexandright_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.
- sampled_div(left: Tensor, right: Tensor, left_index: Optional[Tensor] = None, right_index: Optional[Tensor] = None) Tensor[source]
Performs a sampled division of
leftbyrightaccording to the indices specified inleft_indexandright_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.
- index_sort(inputs: Tensor, max_value: Optional[int] = None) Tuple[Tensor, Tensor][source]
Sorts the elements of the
inputstensor in ascending order. It is expected thatinputsis one-dimensional and that it only contains positive integer values. Ifmax_valueis 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.
- 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 dimensiondim, based on the indices specified viaptr, 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]])
- spline_basis(pseudo: Tensor, kernel_size: Tensor, is_open_spline: Tensor, degree: int = 1) Tuple[Tensor, Tensor][source]
Computes the B-spline basis functions.
- spline_weighting(x: Tensor, weight: Tensor, basis: Tensor, weight_index: Tensor) Tensor[source]
Computes the spline weighting of input features.
- grid_cluster(pos: Tensor, size: Tensor, start: Optional[Tensor] = None, end: Optional[Tensor] = None) Tensor[source]
Clusters all points in
posinto voxels of sizesize.Each point is assigned a cluster index based on which voxel it falls into. The voxel grid is defined by the
sizeparameter and optionally bounded bystartandend.- 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]. IfNone, uses the minimum ofpos.end (
Optional[Tensor], default:None) – End of the voxel grid in each dimension of shape[D]. IfNone, uses the maximum ofpos.
- 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:
- 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
ytheknearest points inx.- 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 forxas a CSR pointer.ptr_y (
Optional[Tensor], default:None) – Batch boundaries foryas a CSR pointer.cosine (
bool, default:False) – IfTrue, 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
xwithin distancerof points iny.- 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 forxas a CSR pointer.ptr_y (
Optional[Tensor], default:None) – Batch boundaries foryas 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) – IfTrue, 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
yfor each point inx.- Parameters:
- Returns:
Tensor– Index tensor of shape[N]with the index of the nearest point inyfor each point inx.
- 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.
- 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
countedges. Ifcount < 1, samplesceil(factor * degree)edges instead.- Parameters:
- Returns:
Tensor– Sampled edge indices (into the edge list).