Add function for the addition of two matrices (rusty1s#177)

* Create spadd.py

Hi,
Maybe it's trivial to have this function, but I still think it'll be helpful and it looks neat when applying matrix addition, i.e., C = A + B.
Thanks

* update

* update

* fix jit

Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>

test/test_add.py

@@ -0,0 +1,33 @@
+from itertools import product
+
+import pytest
+import torch
+from torch_sparse import SparseTensor, add
+
+from .utils import dtypes, devices, tensor
+
+
+@pytest.mark.parametrize('dtype,device', product(dtypes, devices))
+def test_add(dtype, device):
+    rowA = torch.tensor([0, 0, 1, 2, 2], device=device)
+    colA = torch.tensor([0, 2, 1, 0, 1], device=device)
+    valueA = tensor([1, 2, 4, 1, 3], dtype, device)
+    A = SparseTensor(row=rowA, col=colA, value=valueA)
+
+    rowB = torch.tensor([0, 0, 1, 2, 2], device=device)
+    colB = torch.tensor([1, 2, 2, 1, 2], device=device)
+    valueB = tensor([2, 3, 1, 2, 4], dtype, device)
+    B = SparseTensor(row=rowB, col=colB, value=valueB)
+
+    C = A + B
+    rowC, colC, valueC = C.coo()
+
+    assert rowC.tolist() == [0, 0, 0, 1, 1, 2, 2, 2]
+    assert colC.tolist() == [0, 1, 2, 1, 2, 0, 1, 2]
+    assert valueC.tolist() == [1, 2, 5, 4, 1, 1, 5, 4]
+
+    @torch.jit.script
+    def jit_add(A: SparseTensor, B: SparseTensor) -> SparseTensor:
+        return add(A, B)
+
+    jit_add(A, B)

torch_sparse/__init__.py

@@ -65,6 +65,7 @@
 from .eye import eye  # noqa
 from .spmm import spmm  # noqa
 from .spspmm import spspmm  # noqa
+from .spadd import spadd  # noqa
 
 __all__ = [
     'SparseStorage',
@@ -111,5 +112,6 @@
     'eye',
     'spmm',
     'spspmm',
+    'spadd',
     '__version__',
 ]

torch_sparse/add.py

@@ -1,34 +1,69 @@
 from typing import Optional
 
 import torch
+from torch import Tensor
 from torch_scatter import gather_csr
 from torch_sparse.tensor import SparseTensor
 
 
-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...
-        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:
-        raise ValueError(
-            f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or '
-            f'(1, {src.size(1)}, ...), but got size {other.size()}.')
-    if value is not None:
-        value = other.to(value.dtype).add_(value)
+@torch.jit._overload  # noqa: F811
+def add(src, other):  # noqa: F811
+    # type: (SparseTensor, Tensor) -> SparseTensor
+    pass
+
+
+@torch.jit._overload  # noqa: F811
+def add(src, other):  # noqa: F811
+    # type: (SparseTensor, SparseTensor) -> SparseTensor
+    pass
+
+
+def add(src, other):  # noqa: F811
+    if isinstance(other, Tensor):
+        rowptr, col, value = src.csr()
+        if other.size(0) == src.size(0) and other.size(1) == 1:  # Row-wise.
+            other = gather_csr(other.squeeze(1), rowptr)
+        elif other.size(0) == 1 and other.size(1) == src.size(1):  # Col-wise.
+            other = other.squeeze(0)[col]
+        else:
+            raise ValueError(
+                f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or '
+                f'(1, {src.size(1)}, ...), but got size {other.size()}.')
+        if value is not None:
+            value = other.to(value.dtype).add_(value)
+        else:
+            value = other.add_(1)
+        return src.set_value(value, layout='coo')
+
+    elif isinstance(other, SparseTensor):
+        rowA, colA, valueA = src.coo()
+        rowB, colB, valueB = other.coo()
+
+        row = torch.cat([rowA, rowB], dim=0)
+        col = torch.cat([colA, colB], dim=0)
+
+        value: Optional[Tensor] = None
+        if valueA is not None and valueB is not None:
+            value = torch.cat([valueA, valueB], dim=0)
+
+        M = max(src.size(0), other.size(0))
+        N = max(src.size(1), other.size(1))
+        sparse_sizes = (M, N)
+
+        out = SparseTensor(row=row, col=col, value=value,
+                           sparse_sizes=sparse_sizes)
+        out = out.coalesce(reduce='sum')
+        return out
+
     else:
-        value = other.add_(1)
-    return src.set_value(value, layout='coo')
+        raise NotImplementedError
 
 
 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...
+    if other.size(0) == src.size(0) and other.size(1) == 1:  # Row-wise.
         other = gather_csr(other.squeeze(1), rowptr)
-        pass
-    elif other.size(0) == 1 and other.size(1) == src.size(1):  # Col-wise...
+    elif other.size(0) == 1 and other.size(1) == src.size(1):  # Col-wise.
         other = other.squeeze(0)[col]
     else:
         raise ValueError(

torch_sparse/spadd.py

@@ -0,0 +1,18 @@
+import torch
+from torch_sparse import coalesce
+
+
+def spadd(indexA, valueA, indexB, valueB, m, n):
+    """Matrix addition of two sparse matrices.
+
+    Args:
+        indexA (:class:`LongTensor`): The index tensor of first sparse matrix.
+        valueA (:class:`Tensor`): The value tensor of first sparse matrix.
+        indexB (:class:`LongTensor`): The index tensor of second sparse matrix.
+        valueB (:class:`Tensor`): The value tensor of second sparse matrix.
+        m (int): The first dimension of the sparse matrices.
+        n (int): The second dimension of the sparse matrices.
+    """
+    index = torch.cat([indexA, indexB], dim=-1)
+    value = torch.cat([valueA, valueB], dim=0)
+    return coalesce(index=index, value=value, m=m, n=n, op='add')

torch_sparse/storage.py

@@ -292,7 +292,7 @@ def sparse_reshape(self, num_rows: int, num_cols: int):
 
         idx = self.sparse_size(1) * self.row() + self.col()
 
-        row = idx // num_cols
+        row = torch.div(idx, num_cols, rounding_mode='floor')
         col = idx % num_cols
         assert row.dtype == torch.long and col.dtype == torch.long