spmm.py 843 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
from torch import Tensor
rusty1s's avatar
rusty1s committed
2
3
4
from torch_scatter import scatter_add


rusty1s's avatar
rusty1s committed
5
6
def spmm(index: Tensor, value: Tensor, m: int, n: int,
         matrix: Tensor) -> Tensor:
rusty1s's avatar
docs  
rusty1s committed
7
8
9
10
11
    """Matrix product of sparse matrix with dense matrix.

    Args:
        index (:class:`LongTensor`): The index tensor of sparse matrix.
        value (:class:`Tensor`): The value tensor of sparse matrix.
wang-ps's avatar
wang-ps committed
12
13
        m (int): The first dimension of sparse matrix.
        n (int): The second dimension of sparse matrix.
rusty1s's avatar
docs  
rusty1s committed
14
15
16
17
        matrix (:class:`Tensor`): The dense matrix.

    :rtype: :class:`Tensor`
    """
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
cleanup  
rusty1s committed
19
    assert n == matrix.size(-2)
rusty1s's avatar
rusty1s committed
20

rusty1s's avatar
rusty1s committed
21
    row, col = index[0], index[1]
rusty1s's avatar
rusty1s committed
22
23
    matrix = matrix if matrix.dim() > 1 else matrix.unsqueeze(-1)

rusty1s's avatar
cleanup  
rusty1s committed
24
    out = matrix.index_select(-2, col)
rusty1s's avatar
rusty1s committed
25
    out = out * value.unsqueeze(-1)
26
    out = scatter_add(out, row, dim=-2, dim_size=m)
rusty1s's avatar
rusty1s committed
27
28

    return out