pyg_lib.ops

grouped_matmul(inputs: List[Tensor], others: List[Tensor]) → List[Tensor]

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[torch.Tensor]) – List of left operand 2D matrices of shapes [N_i, K_i].

  • others (List[torch.Tensor]) – List of right operand 2D matrices of shapes [K_i, M_i].

Returns

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

Return type

List[torch.Tensor]

segment_matmul(inputs: Tensor, ptr: Tensor, other: Tensor) → Tensor

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

  • input (torch.Tensor) – The left operand 2D matrix of shape [N, K].

  • ptr (torch.Tensor) – Compressed vector of shape [B + 1], holding the boundaries of segments. For best performance, given as a CPU tensor.

  • other (torch.Tensor) – The right operand 3D tensor of shape [B, K, M].

Returns

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

Return type

torch.Tensor