spcadd, sparseMask, cadd, csub, cmul + tests

Author: Martin Raison

test/common.py

@@ -128,11 +128,12 @@ def assertEqual(self, x, y, prec=None, message=''):
 
         if torch.is_tensor(x) and torch.is_tensor(y):
             def assertTensorsEqual(a, b):
-                max_err = 0
                 super(TestCase, self).assertEqual(a.size(), b.size())
-                for index in iter_indices(a):
-                    max_err = max(max_err, abs(a[index] - b[index]))
-                self.assertLessEqual(max_err, prec, message)
+                if a.numel() > 0:
+                    b = b.type_as(a)
+                    b = b.cuda(device=a.get_device()) if a.is_cuda else b.cpu()
+                    max_err = (a - b).abs().max()
+                    self.assertLessEqual(max_err, prec, message)
             self.assertEqual(x.is_sparse, y.is_sparse, message)
             if x.is_sparse:
                 assertTensorsEqual(x.indices(), y.indices())
@@ -161,11 +162,12 @@ def assertNotEqual(self, x, y, prec=None, message=''):
             y = y.data
 
         if torch.is_tensor(x) and torch.is_tensor(y):
-            max_err = 0
             if x.size() != y.size():
                 super(TestCase, self).assertNotEqual(x.size(), y.size())
-            for index in iter_indices(x):
-                max_err = max(max_err, abs(x[index] - y[index]))
+            self.assertGreater(x.numel(), 0)
+            y = y.type_as(x)
+            y = y.cuda(device=x.get_device()) if x.is_cuda else y.cpu()
+            max_err = (x - y).abs().max()
             self.assertGreaterEqual(max_err, prec, message)
         elif type(x) == str and type(y) == str:
             super(TestCase, self).assertNotEqual(x, y)

test/test_sparse.py

@@ -15,7 +15,7 @@
 type_triplets = [cpu_triplet]
 if torch.cuda.is_available():
     cuda_triplet = (
-        torch.cuda.LongTensor,
+        torch.cuda.IntTensor,
         torch.cuda.DoubleTensor,
         torch.cuda.sparse.DoubleTensor)
     type_triplets.append(cuda_triplet)
@@ -312,27 +312,33 @@ def test_shape(di, dj, dk):
         test_shape(3000, 64, 300)
 
     def _test_spadd_shape(self, shape_i, shape_v=None):
-        shape = shape_i + (shape_v or [])
-        x, _, _ = self._gen_sparse(len(shape_i), 10, shape)
-        y = torch.randn(*shape)
-        r = random.random()
+        for is_cuda in [False, True]:
+            shape = shape_i + (shape_v or [])
+            x, _, _ = self._gen_sparse(len(shape_i), 10, shape, is_cuda)
+            y = torch.randn(*shape)
+            if is_cuda:
+                y = y.cuda()
+            r = random.random()
 
-        expected = y + r * x.to_dense()
-        res = torch.add(y, r, x)
+            expected = y + r * x.to_dense()
+            res = torch.add(y, r, x)
 
-        self.assertEqual(res, expected)
+            self.assertEqual(res, expected)
 
-        # Non contiguous dense tensor
-        s = list(shape)
-        s[0] = shape[-1]
-        s[-1] = shape[0]
-        y = torch.randn(*s).transpose_(0, len(s) - 1)
-        r = random.random()
+            # Non contiguous dense tensor
+            s = list(shape)
+            s[0] = shape[-1]
+            s[-1] = shape[0]
+            y = torch.randn(*s)
+            if is_cuda:
+                y = y.cuda()
+            y.transpose_(0, len(s) - 1)
+            r = random.random()
 
-        expected = y + r * x.to_dense()
-        res = torch.add(y, r, x)
+            expected = y + r * x.to_dense()
+            res = torch.add(y, r, x)
 
-        self.assertEqual(res, expected)
+            self.assertEqual(res, expected)
 
     def test_spadd(self):
         self._test_spadd_shape([5, 6])
@@ -347,49 +353,50 @@ def test_spadd_hybrid(self):
         self._test_spadd_shape([5, 5, 5, 5, 5, 5], [2])
 
     def _test_basic_ops_shape(self, shape_i, shape_v=None):
-        shape = shape_i + (shape_v or [])
-        x1, _, _ = self._gen_sparse(len(shape_i), 9, shape)
-        x2, _, _ = self._gen_sparse(len(shape_i), 12, shape)
-
-        y1 = x1 + x2
-        y2 = x1.clone()
-        y2.add_(x2)
-        expected = x1.to_dense() + x2.to_dense()
-        self.assertEqual(y1.to_dense(), expected)
-        self.assertEqual(y2.to_dense(), expected)
-
-        y1 = x1 - x2
-        y2 = x1.clone()
-        y2.sub_(x2)
-        expected = x1.to_dense() - x2.to_dense()
-        self.assertEqual(y1.to_dense(), expected)
-        self.assertEqual(y2.to_dense(), expected)
-
-        y1 = x1 * x2
-        y2 = x1.clone()
-        y2.mul_(x2)
-        expected = x1.to_dense() * x2.to_dense()
-        self.assertEqual(y1.to_dense(), expected)
-        self.assertEqual(y2.to_dense(), expected)
-
-        y1 = x1 * 37.5
-        y2 = x1.clone()
-        y2.mul_(37.5)
-        expected = x1.to_dense() * 37.5
-        self.assertEqual(y1.to_dense(), expected)
-        self.assertEqual(y2.to_dense(), expected)
-
-        y1 = x1 / 37.5
-        y2 = x1.clone()
-        y2.div_(37.5)
-        expected = x1.to_dense() / 37.5
-        self.assertEqual(y1.to_dense(), expected)
-        self.assertEqual(y2.to_dense(), expected)
-
-        y = x1.clone()
-        y.zero_()
-        expected = torch.zeros(x1.size())
-        self.assertEqual(y.to_dense(), expected)
+        for is_cuda in [False, True]:
+            shape = shape_i + (shape_v or [])
+            x1, _, _ = self._gen_sparse(len(shape_i), 9, shape, is_cuda)
+            x2, _, _ = self._gen_sparse(len(shape_i), 12, shape, is_cuda)
+
+            y1 = x1 + x2
+            y2 = x1.clone()
+            y2.add_(x2)
+            expected = x1.to_dense() + x2.to_dense()
+            self.assertEqual(y1.to_dense(), expected)
+            self.assertEqual(y2.to_dense(), expected)
+
+            y1 = x1 - x2
+            y2 = x1.clone()
+            y2.sub_(x2)
+            expected = x1.to_dense() - x2.to_dense()
+            self.assertEqual(y1.to_dense(), expected)
+            self.assertEqual(y2.to_dense(), expected)
+
+            y1 = x1 * x2
+            y2 = x1.clone()
+            y2.mul_(x2)
+            expected = x1.to_dense() * x2.to_dense()
+            self.assertEqual(y1.to_dense(), expected)
+            self.assertEqual(y2.to_dense(), expected)
+
+            y1 = x1 * 37.5
+            y2 = x1.clone()
+            y2.mul_(37.5)
+            expected = x1.to_dense() * 37.5
+            self.assertEqual(y1.to_dense(), expected)
+            self.assertEqual(y2.to_dense(), expected)
+
+            y1 = x1 / 37.5
+            y2 = x1.clone()
+            y2.div_(37.5)
+            expected = x1.to_dense() / 37.5
+            self.assertEqual(y1.to_dense(), expected)
+            self.assertEqual(y2.to_dense(), expected)
+
+            y = x1.clone()
+            y.zero_()
+            expected = torch.zeros(x1.size())
+            self.assertEqual(y.to_dense(), expected)
 
     def test_basic_ops(self):
         self._test_basic_ops_shape([5, 6])
@@ -403,6 +410,59 @@ def test_basic_ops_hybrid(self):
         self._test_basic_ops_shape([50, 30, 20], [2])
         self._test_basic_ops_shape([5, 5, 5, 5, 5, 5], [2])
 
+    def _test_sparse_mask_shape(self, shape_i, shape_v=None):
+        for is_cuda in [False, True]:
+            shape = shape_i + (shape_v or [])
+            x1, _, _ = self._gen_sparse(len(shape_i), 9, shape, is_cuda)
+            x2, _, _ = self._gen_sparse(len(shape_i), 12, shape, is_cuda)
+
+            y1 = x1 + x2
+            y2 = x1.clone()
+            y2.add_(x2)
+            expected = x1.to_dense() + x2.to_dense()
+            self.assertEqual(y1.to_dense(), expected)
+            self.assertEqual(y2.to_dense(), expected)
+
+    def test_sparse_mask(self):
+        for IndexTensor, ValueTensor, SparseTensor in type_triplets:
+            i = IndexTensor([
+                [1, 3, 3, 0, 4],
+                [2, 1, 1, 2, 3],
+            ])
+            v = ValueTensor([1, 2, 3, 4, 5])
+            x = SparseTensor(i, v, torch.Size([5, 4]))
+            dense = ValueTensor([
+                [1, 2, 3, 4],
+                [5, 6, 7, 8],
+                [9, 10, 11, 12],
+                [13, 14, 15, 16],
+                [17, 18, 19, 20],
+            ])
+            exp_v = ValueTensor([7, 14, 14, 3, 20])
+            expected = SparseTensor(i, exp_v, torch.Size([5, 4]))
+            res = dense.sparse_mask(x)
+            self.assertEqual(res, expected)
+
+    def test_sparse_mask_hybrid(self):
+        for IndexTensor, ValueTensor, SparseTensor in type_triplets:
+            i = IndexTensor([
+                [1, 3, 3, 0, 4],
+                [2, 1, 1, 2, 3],
+            ])
+            v = ValueTensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6]])
+            x = SparseTensor(i, v, torch.Size([5, 4, 2]))
+            dense = ValueTensor([
+                [[1, 3], [2, 2], [3, 3], [4, 2]],
+                [[5, 7], [6, 7], [7, 9], [8, 9]],
+                [[9, 2], [10, 4], [11, 1], [12, 3]],
+                [[13, 5], [14, 1], [15, 1], [16, 6]],
+                [[17, 7], [18, 2], [19, 7], [20, 1]],
+            ])
+            exp_v = ValueTensor([[7, 9], [14, 1], [14, 1], [3, 3], [20, 1]])
+            expected = SparseTensor(i, exp_v, torch.Size([5, 4, 2]))
+            res = dense.sparse_mask(x)
+            self.assertEqual(res, expected)
+
 
 if __name__ == '__main__':
     run_tests()

tools/cwrap/plugins/THPPlugin.py

@@ -19,6 +19,7 @@ class THPPlugin(CWrapPlugin):
 
         'THCudaTensor*': Template('((THCPFloatTensor*)$arg)->cdata'),
         'THCudaDoubleTensor*': Template('((THCPDoubleTensor*)$arg)->cdata'),
+        'THCudaIntTensor*': Template('((THCPIntTensor*)$arg)->cdata'),
         'THCudaLongTensor*': Template('((THCPLongTensor*)$arg)->cdata'),
 
         'THSFloatTensor*': Template('((THSPFloatTensor*)$arg)->cdata'),
@@ -56,6 +57,7 @@ class THPPlugin(CWrapPlugin):
 
         'THCudaTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPFloatTensorClass'),
         'THCudaDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPDoubleTensorClass'),
+        'THCudaIntTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPIntTensorClass'),
         'THCudaLongTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPLongTensorClass'),
 
         'THSDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPDoubleTensorClass'),
@@ -86,8 +88,10 @@ class THPPlugin(CWrapPlugin):
     RETURN_WRAPPER = {
         'THTensor*': Template('return THPTensor_(New)($result);'),
         'THSTensor*': Template('return THSPTensor_(New)($result);'),
+        'THIndexTensor*': Template('return THPIndexTensor_(New)($result);'),
         'THLongTensor*': Template('return THPLongTensor_New($result);'),
         'THLongStorage*': Template('return THPLongStorage_New($result);'),
+        'THCudaIntTensor*': Template('return THCPIntTensor_New($result);'),
         'THCudaLongTensor*': Template('return THCPLongTensor_New($result);'),
         # TODO: make it smarter - it should return python long if result doesn't fit into an int
         'long': Template('return PyInt_FromLong($result);'),
@@ -174,6 +178,7 @@ def _allocate(typename, tmpl, cuda_tmpl=None, sparse=False):
         'THDoubleTensor*': '" THPModuleStr "DoubleTensor',
         'THCudaTensor*': 'torch.cuda.FloatTensor',
         'THCudaDoubleTensor*': 'torch.cuda.DoubleTensor',
+        'THCudaIntTensor*': 'torch.cuda.IntTensor',
         'THCudaLongTensor*': 'torch.cuda.LongTensor',
         'THSize*': 'torch.Size',
         'THStride*': 'tuple',

torch/_utils.py

@@ -57,7 +57,7 @@ def _cuda(self, device=None, async=False):
     with torch.cuda.device(device):
         if self.is_sparse:
             new_type = getattr(torch.cuda.sparse, self.__class__.__name__)
-            indices = self.indices().cuda(device, async)
+            indices = self.indices().cuda(device, async).int()
             values = self.values().cuda(device, async)
             return new_type(indices, values, self.size())
         else:

torch/csrc/generic/SparseTensor.cpp

@@ -26,9 +26,9 @@ static void THSPTensor_(dealloc)(THSPTensor* self)
 static PyObject * THSPTensor_(pynew)(PyTypeObject *type, PyObject *args, PyObject *kwargs)
 {
 #ifdef THC_GENERIC_FILE
-#define THPIndexTensor_Check THCPLongTensor_Check
-#define THPIndexTensor THCPLongTensor
-#define THIndexTensor THCudaLongTensor
+#define THPIndexTensor_Check THCPIntTensor_Check
+#define THPIndexTensor THCPIntTensor
+#define THIndexTensor THCudaIntTensor
 #else
 #define THPIndexTensor_Check THPLongTensor_Check
 #define THPIndexTensor THPLongTensor

torch/csrc/generic/methods/SparseTensor.cwrap

@@ -83,12 +83,11 @@ PyObject * THSPTensor_(size)(PyObject *self, PyObject *args, PyObject *kwargs)
   name: indices
   defined_if: "IS_CUDA"
   sparse: yes
-  return: THCudaLongTensor*
+  return: THCudaIntTensor*
   arguments:
   - THSTensor* self
 ]]
 
-
 [[
   name: values
   sparse: yes
@@ -372,3 +371,13 @@ PyObject * THSPTensor_(size)(PyObject *self, PyObject *args, PyObject *kwargs)
     - THTensor* self
     - THSTensor* mask
 ]]
+
+[[
+  name: getDevice
+  sparse: yes
+  python_name: get_device
+  defined_if: IS_CUDA
+  return: long
+  arguments:
+    - THSTensor* self
+]]