dgl.sparse.bspmm

dgl.sparse.bspmm(A: SparseMatrix, X: Tensor) Tensor[源码]

按批次将稀疏矩阵与稠密矩阵相乘,等同于 A @ X

参数:
  • A (SparseMatrix) – 形状为 (L, M) 的稀疏矩阵,其向量值的长度为 K

  • X (torch.Tensor) – 形状为 (M, N, K) 的稠密矩阵

返回:

形状为 (L, N, K) 的稠密矩阵

返回类型:

torch.Tensor

示例

>>> indices = torch.tensor([[0, 1, 1], [1, 0, 2]])
>>> val = torch.randn(len(row), 2)
>>> A = dglsp.spmatrix(indices, val, shape=(3, 3))
>>> X = torch.randn(3, 3, 2)
>>> result = dglsp.bspmm(A, X)
>>> type(result)
<class 'torch.Tensor'>
>>> result.shape
torch.Size([3, 3, 2])