add.py 2.5 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
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...
rusty1s's avatar
rusty1s committed
11
        other = gather_csr(other.squeeze(1), rowptr)
rusty1s's avatar
rusty1s committed
12
        pass
rusty1s's avatar
rusty1s committed
13
14
    elif other.size(0) == 1 and other.size(1) == src.size(1):  # Col-wise...
        other = other.squeeze(0)[col]
rusty1s's avatar
rusty1s committed
15
    else:
rusty1s's avatar
rusty1s committed
16
        raise ValueError(
rusty1s's avatar
rusty1s committed
17
18
            f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or '
            f'(1, {src.size(1)}, ...), but got size {other.size()}.')
rusty1s's avatar
rusty1s committed
19
    if value is not None:
rusty1s's avatar
fixes  
rusty1s committed
20
        value = other.to(value.dtype).add_(value)
rusty1s's avatar
rusty1s committed
21
22
23
    else:
        value = other.add_(1)
    return src.set_value(value, layout='coo')
rusty1s's avatar
rusty1s committed
24
25


rusty1s's avatar
rusty1s committed
26
27
28
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...
rusty1s's avatar
rusty1s committed
29
        other = gather_csr(other.squeeze(1), rowptr)
rusty1s's avatar
rusty1s committed
30
31
32
33
        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
34
        raise ValueError(
rusty1s's avatar
rusty1s committed
35
36
            f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or '
            f'(1, {src.size(1)}, ...), but got size {other.size()}.')
rusty1s's avatar
rusty1s committed
37

rusty1s's avatar
rusty1s committed
38
    if value is not None:
rusty1s's avatar
fixes  
rusty1s committed
39
        value = value.add_(other.to(value.dtype))
rusty1s's avatar
rusty1s committed
40
41
42
    else:
        value = other.add_(1)
    return src.set_value_(value, layout='coo')
rusty1s's avatar
rusty1s committed
43

rusty1s's avatar
rusty1s committed
44

rusty1s's avatar
rusty1s committed
45
46
47
48
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
49
        value = value.add(other.to(value.dtype))
rusty1s's avatar
rusty1s committed
50
51
52
    else:
        value = other.add(1)
    return src.set_value(value, layout=layout)
rusty1s's avatar
rusty1s committed
53
54


rusty1s's avatar
rusty1s committed
55
56
57
58
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
59
        value = value.add_(other.to(value.dtype))
rusty1s's avatar
rusty1s committed
60
61
62
    else:
        value = other.add(1)
    return src.set_value_(value, layout=layout)
rusty1s's avatar
rusty1s committed
63
64


rusty1s's avatar
rusty1s committed
65
66
67
68
69
70
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
71
72
73
SparseTensor.__add__ = SparseTensor.add
SparseTensor.__radd__ = SparseTensor.add
SparseTensor.__iadd__ = SparseTensor.add_