add.py 2.66 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
from typing import Optional

rusty1s's avatar
rusty1s committed
3
import torch
rusty1s's avatar
rusty1s committed
4
from torch_scatter import gather_csr
rusty1s's avatar
rusty1s committed
5
from torch_sparse.tensor import SparseTensor
rusty1s's avatar
rusty1s committed
6

rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
9
10
11
12
13
@torch.jit.script
def add(src: SparseTensor, other: torch.Tensor) -> SparseTensor:
    rowptr, col, value = src.csr()
    if other.size(0) == src.size(0) and other.size(1) == 1:  # Row-wise...
        # TODO
        # other = gather_csr(other.squeeze(1), rowptr)
rusty1s's avatar
rusty1s committed
14
        pass
rusty1s's avatar
rusty1s committed
15
16
    elif other.size(0) == 1 and other.size(1) == src.size(1):  # Col-wise...
        other = other.squeeze(0)[col]
rusty1s's avatar
rusty1s committed
17
18
19
20
    else:
        raise ValueError(f'Size mismatch: Expected size ({src.size(0)}, 1,'
                         f' ...) or (1, {src.size(1)}, ...), but got size '
                         f'{other.size()}.')
rusty1s's avatar
rusty1s committed
21

rusty1s's avatar
rusty1s committed
22
    if value is not None:
rusty1s's avatar
fixes  
rusty1s committed
23
        value = other.to(value.dtype).add_(value)
rusty1s's avatar
rusty1s committed
24
25
26
    else:
        value = other.add_(1)
    return src.set_value(value, layout='coo')
rusty1s's avatar
rusty1s committed
27
28


rusty1s's avatar
rusty1s committed
29
30
31
32
33
34
35
36
37
38
@torch.jit.script
def add_(src: SparseTensor, other: torch.Tensor) -> SparseTensor:
    rowptr, col, value = src.csr()
    if other.size(0) == src.size(0) and other.size(1) == 1:  # Row-wise...
        # TODO
        # other = gather_csr(other.squeeze(1), rowptr)
        pass
    elif other.size(0) == 1 and other.size(1) == src.size(1):  # Col-wise...
        other = other.squeeze(0)[col]
    else:
rusty1s's avatar
rusty1s committed
39
40
41
42
        raise ValueError(f'Size mismatch: Expected size ({src.size(0)}, 1,'
                         f' ...) or (1, {src.size(1)}, ...), but got size '
                         f'{other.size()}.')

rusty1s's avatar
rusty1s committed
43
    if value is not None:
rusty1s's avatar
fixes  
rusty1s committed
44
        value = value.add_(other.to(value.dtype))
rusty1s's avatar
rusty1s committed
45
46
47
    else:
        value = other.add_(1)
    return src.set_value_(value, layout='coo')
rusty1s's avatar
rusty1s committed
48

rusty1s's avatar
rusty1s committed
49

rusty1s's avatar
rusty1s committed
50
51
52
53
54
@torch.jit.script
def add_nnz(src: SparseTensor, other: torch.Tensor,
            layout: Optional[str] = None) -> SparseTensor:
    value = src.storage.value()
    if value is not None:
rusty1s's avatar
fixes  
rusty1s committed
55
        value = value.add(other.to(value.dtype))
rusty1s's avatar
rusty1s committed
56
57
58
    else:
        value = other.add(1)
    return src.set_value(value, layout=layout)
rusty1s's avatar
rusty1s committed
59
60


rusty1s's avatar
rusty1s committed
61
62
63
64
65
@torch.jit.script
def add_nnz_(src: SparseTensor, other: torch.Tensor,
             layout: Optional[str] = None) -> SparseTensor:
    value = src.storage.value()
    if value is not None:
rusty1s's avatar
fixes  
rusty1s committed
66
        value = value.add_(other.to(value.dtype))
rusty1s's avatar
rusty1s committed
67
68
69
    else:
        value = other.add(1)
    return src.set_value_(value, layout=layout)
rusty1s's avatar
rusty1s committed
70
71


rusty1s's avatar
rusty1s committed
72
73
74
75
76
77
SparseTensor.add = lambda self, other: add(self, other)
SparseTensor.add_ = lambda self, other: add_(self, other)
SparseTensor.add_nnz = lambda self, other, layout=None: add_nnz(
    self, other, layout)
SparseTensor.add_nnz_ = lambda self, other, layout=None: add_nnz_(
    self, other, layout)
rusty1s's avatar
fixes  
rusty1s committed
78
79
80
SparseTensor.__add__ = SparseTensor.add
SparseTensor.__radd__ = SparseTensor.add
SparseTensor.__iadd__ = SparseTensor.add_