transpose.py 1.86 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
import torch

rusty1s's avatar
rusty1s committed
3
4
5
from torch_sparse.storage import SparseStorage
from torch_sparse.tensor import SparseTensor

rusty1s's avatar
rusty1s committed
6

rusty1s's avatar
rusty1s committed
7
@torch.jit.script
rusty1s's avatar
rusty1s committed
8
def t(src: SparseTensor) -> SparseTensor:
rusty1s's avatar
rusty1s committed
9
10
11
12
13
14
    csr2csc = src.storage.csr2csc()

    row, col, value = src.coo()

    if value is not None:
        value = value[csr2csc]
rusty1s's avatar
rusty1s committed
15

rusty1s's avatar
rusty1s committed
16
17
18
19
    sparse_sizes = src.storage.sparse_sizes()

    storage = SparseStorage(
        row=col[csr2csc],
rusty1s's avatar
rusty1s committed
20
        rowptr=src.storage._colptr,
rusty1s's avatar
rusty1s committed
21
22
23
        col=row[csr2csc],
        value=value,
        sparse_sizes=torch.Size([sparse_sizes[1], sparse_sizes[0]]),
rusty1s's avatar
rusty1s committed
24
25
        rowcount=src.storage._colcount,
        colptr=src.storage._rowptr,
rusty1s's avatar
rusty1s committed
26
        colcount=src.storage._rowcount,
rusty1s's avatar
rusty1s committed
27
28
29
30
31
32
        csr2csc=src.storage._csc2csr,
        csc2csr=csr2csc,
        is_sorted=True,
    )

    return src.from_storage(storage)
rusty1s's avatar
rusty1s committed
33
34
35


SparseTensor.t = lambda self: t(self)
rusty1s's avatar
rusty1s committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

###############################################################################


def transpose(index, value, m, n, coalesced=True):
    """Transposes dimensions 0 and 1 of a sparse tensor.

    Args:
        index (:class:`LongTensor`): The index tensor of sparse matrix.
        value (:class:`Tensor`): The value tensor of sparse matrix.
        m (int): The first dimension of corresponding dense matrix.
        n (int): The second dimension of corresponding dense matrix.
        coalesced (bool, optional): If set to :obj:`False`, will not coalesce
            the output. (default: :obj:`True`)
    :rtype: (:class:`LongTensor`, :class:`Tensor`)
    """

    row, col = index
rusty1s's avatar
rusty1s committed
54
55
    row, col = col, row

rusty1s's avatar
rusty1s committed
56
    if coalesced:
rusty1s's avatar
rusty1s committed
57
58
59
60
61
62
63
        sparse_sizes = torch.Size([n, m])
        storage = SparseStorage(row=row, col=col, value=value,
                                sparse_sizes=sparse_sizes, is_sorted=False)
        storage = storage.coalesce()
        row, col, value = storage.row(), storage.col(), storage.value()

    return torch.stack([row, col], dim=0), value