dgl.ops.segment_mm

dgl.ops.segment_mm(a, b, seglen_a)[source]

Performs matrix multiplication according to segments.

Suppose seglen_a == [10, 5, 0, 3], the operator will perform four matrix multiplications:

a[0:10] @ b[0], a[10:15] @ b[1],
a[15:15] @ b[2], a[15:18] @ b[3]
Parameters

  • a (Tensor) – The left operand, 2-D tensor of shape (N, D1)

  • b (Tensor) – The right operand, 3-D tensor of shape (R, D1, D2)

  • seglen_a (Tensor) – An integer tensor of shape (R,). Each element is the length of segments of input a. The summation of all elements must be equal to N.

Returns

The output dense matrix of shape (N, D2)

Return type

Tensor