Better support for adding zero-filled sparse tensors (#7479)

zou3519 · web-flow · commit 56e7a2cde10b · 2018-05-18T10:29:27.000-04:00
Right now, if we add a zero-filled sparse tensor with another sparse
tensor, both tensors must have the same "density" (dimI, dimV) and size
(tensor.size()) for them to be added successfully. This relaxes that
constraint so that if both tensors have the same tensor.size() and at
least one is zero-filled, they can be added successfully.

Before:
```
i = torch.LongTensor([[0, 1, 1], [2, 0, 2]])
v = torch.FloatTensor([3, 4, 5]).unsqueeze(1)
sparse_mat = torch.sparse.FloatTensor(i, v, torch.Size([2,3,1]))
zeros = torch.zeros(sparse_mat.size(), layout=torch.sparse_coo)
sparse_mat + zeros

RuntimeError: cadd operands have
incompatible sizes or dimension types
at
../src/THS/generic/THSTensorMath.c:126
```

After: no error.
diff --git a/aten/src/THCS/generic/THCSTensorMath.cu b/aten/src/THCS/generic/THCSTensorMath.cu
@@ -367,10 +367,29 @@ void THCSTensor_(div)(THCState *state, THCSTensor *r_, THCSTensor *t, real value
   }
 }
 
+int THCSTensor_(isSameSizeIgnoringDensity)(THCState *state, const THCSTensor *self, const THCSTensor *src) {
+  int d;
+  if (self->nDimensionI + self->nDimensionV != src->nDimensionI + src->nDimensionV) {
+    return 0;
+  }
+  for(d = 0; d < self->nDimensionI + self->nDimensionV; ++d) {
+    if(self->size[d] != src->size[d]) {
+      return 0;
+    }
+  }
+  return 1;
+}
+
+int THCSTensor_(isSameDensity)(THCState *state, const THCSTensor *self, const THCSTensor *src) {
+  return self->nDimensionI == src->nDimensionI &&
+      self->nDimensionV == src->nDimensionV;
+}
+
+
 void THCSTensor_(cadd)(THCState *state, THCSTensor *r_, THCSTensor *t, real value, THCSTensor *src) {
   THCAssertSameGPU(THCSTensor_(checkGPU)(state, 3, 3, r_, t, src));
-  if(!THCSTensor_(isSameSizeAs)(state, t, src)) {
-    THError("cadd operands have incompatible sizes or dimension types");
+  if (!THCSTensor_(isSameSizeIgnoringDensity)(state, t, src)) {
+    THError("cadd operands have incompatible sizes");
   }
 
   if (src->nnz == 0) {
@@ -382,6 +401,10 @@ void THCSTensor_(cadd)(THCState *state, THCSTensor *r_, THCSTensor *t, real valu
     return;
   }
 
+  if(!THCSTensor_(isSameDensity)(state, t, src)) {
+    THError("cadd operands have incompatible densities");
+  }
+
   // We deliberately choose to simply concat the indices and values tensors
   // rather than merging them. This removes the need to synchronously fetch nnz
   // at the end of the operation, at the cost of having a non-coalesced result.
diff --git a/aten/src/THS/generic/THSTensorMath.c b/aten/src/THS/generic/THSTensorMath.c
@@ -121,9 +121,27 @@ void THSTensor_(div)(THSTensor *r_, THSTensor *t, real value) {
   }
 }
 
+int THSTensor_(isSameSizeIgnoringDensity)(const THSTensor *self, const THSTensor* src) {
+  int d;
+  if (self->nDimensionI + self->nDimensionV != src->nDimensionI + src->nDimensionV) {
+    return 0;
+  }
+  for(d = 0; d < self->nDimensionI + self->nDimensionV; ++d) {
+    if(self->size[d] != src->size[d]) {
+      return 0;
+    }
+  }
+  return 1;
+}
+
+int THSTensor_(isSameDensity)(const THSTensor *self, const THSTensor* src) {
+  return self->nDimensionI == src->nDimensionI &&
+      self->nDimensionV == src->nDimensionV;
+}
+
 void THSTensor_(cadd)(THSTensor *r_, THSTensor *t, real value, THSTensor *src) {
-  if(!THSTensor_(isSameSizeAs)(t, src)) {
-    THError("cadd operands have incompatible sizes or dimension types");
+  if (!THSTensor_(isSameSizeIgnoringDensity)(t, src)) {
+    THError("cadd operands have incompatible sizes");
   }
 
   if (src->nnz == 0) {
@@ -135,6 +153,10 @@ void THSTensor_(cadd)(THSTensor *r_, THSTensor *t, real value, THSTensor *src) {
     return;
   }
 
+  if(!THSTensor_(isSameDensity)(t, src)) {
+    THError("cadd operands have incompatible densities");
+  }
+
   // saving those because they can be overwritten when doing in-place operations
   ptrdiff_t t_nnz = t->nnz, s_nnz = src->nnz, max_nnz = t_nnz + s_nnz;
   int t_coalesced = t->coalesced, s_coalesced = src->coalesced;
diff --git a/test/test_sparse.py b/test/test_sparse.py
@@ -385,6 +385,17 @@ def test_transpose_coalesce_invariant(self):
         self.assertTrue(x_coalesced.is_coalesced())
         self.assertFalse(y_uncoalesced.is_coalesced())
 
+    def test_add_zeros(self):
+        def test_shape(sparse_dims, sizes):
+            x, _, _ = self._gen_sparse(sparse_dims, 20, sizes)
+            zeros = torch.zeros(sizes, layout=torch.sparse_coo).to(x.device)
+            self.assertEqual(zeros + x, x)
+            self.assertEqual(x + zeros, x)
+
+        test_shape(1, [1])
+        test_shape(4, [3, 17, 19, 5])
+        test_shape(2, [3, 17, 19, 5])
+
     @cpu_only
     def test_mm(self):
         def test_shape(di, dj, dk):
