dgl.sparse.bsddmm

dgl.sparse.bsddmm(A: SparseMatrix, X1: Tensor, X2: Tensor) SparseMatrix[source]

按批次进行采样-密集-密集矩阵乘法(SDDMM)。

sddmm matrix-multiplies two dense matrices X1 and X2, then elementwise-multiplies the result with sparse matrix A at the nonzero locations.

Mathematically sddmm is formulated as:

\[out = (X1 @ X2) * A\]

批次维度是输入密集矩阵的最后一个维度。特别是,如果稀疏矩阵具有标量非零值,它将在bsddmm中进行广播。

Parameters:
  • A (SparseMatrix) – 形状为 (L, N) 的稀疏矩阵,具有标量值或长度为 K 的向量值

  • X1 (张量) – 形状为 (L, M, K) 的密集矩阵

  • X2 (张量) – 形状为 (M, N, K) 的密集矩阵

Returns:

形状为(L, N)的稀疏矩阵,其向量长度为K

Return type:

SparseMatrix

示例

>>> indices = torch.tensor([[1, 1, 2], [2, 3, 3]])
>>> val = torch.arange(1, 4).float()
>>> A = dglsp.spmatrix(indices, val, (3, 4))
>>> X1 = torch.arange(0, 3 * 5 * 2).view(3, 5, 2).float()
>>> X2 = torch.arange(0, 5 * 4 * 2).view(5, 4, 2).float()
>>> dglsp.bsddmm(A, X1, X2)
SparseMatrix(indices=tensor([[1, 1, 2],
                             [2, 3, 3]]),
             values=tensor([[1560., 1735.],
                            [3400., 3770.],
                            [8400., 9105.]]),
             shape=(3, 4), nnz=3, val_size=(2,))